From 3e418d2bcbc2f0e51e75d0e243b58d9e1b4d8805 Mon Sep 17 00:00:00 2001 From: Paul Slaughter Date: Tue, 3 Dec 2024 10:11:38 -0600 Subject: [PATCH 1/2] Feat(Q): Add AmazonQ db and config - Add usage of feature flag to appease pipeline - Remove `.com` check since FF check should be sufficient (and we just won't turn the feature on in `.com`) --- ...41120210519_add_amazon_q_to_ai_settings.rb | 38 ++++++++++++++ db/schema_migrations/20241120210519 | 1 + db/structure.sql | 15 ++++++ ee/app/models/ai/setting.rb | 13 +++++ .../models/gitlab_subscriptions/features.rb | 7 +-- .../wip/amazon_q_integration.yml | 9 ++++ ee/lib/ai/amazon_q.rb | 18 +++++++ ee/spec/lib/ai/amazon_q_spec.rb | 49 +++++++++++++++++++ ee/spec/models/ai/setting_spec.rb | 18 +++++++ 9 files changed, 165 insertions(+), 3 deletions(-) create mode 100644 db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb create mode 100644 db/schema_migrations/20241120210519 create mode 100644 ee/config/feature_flags/wip/amazon_q_integration.yml create mode 100644 ee/lib/ai/amazon_q.rb create mode 100644 ee/spec/lib/ai/amazon_q_spec.rb diff --git a/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb b/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb new file mode 100644 index 00000000000000..1b78acf77b1814 --- /dev/null +++ b/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb @@ -0,0 +1,38 @@ +# frozen_string_literal: true + +class AddAmazonQToAiSettings < Gitlab::Database::Migration[2.2] + disable_ddl_transaction! + + milestone '17.7' + + def up + with_lock_retries do + add_column :ai_settings, :amazon_q_oauth_application_id, :bigint + add_column :ai_settings, :amazon_q_service_account_user_id, :bigint + add_column :ai_settings, :amazon_q_ready, :boolean, default: false, null: false + add_column :ai_settings, :amazon_q_role_arn, :text + end + + add_concurrent_index :ai_settings, :amazon_q_oauth_application_id + add_concurrent_foreign_key :ai_settings, + :oauth_applications, + column: :amazon_q_oauth_application_id, + on_delete: :nullify + add_concurrent_index :ai_settings, :amazon_q_service_account_user_id + add_concurrent_foreign_key :ai_settings, + :users, + column: :amazon_q_service_account_user_id, + on_delete: :nullify + + add_text_limit :ai_settings, :amazon_q_role_arn, 2048 + end + + def down + with_lock_retries do + remove_column :ai_settings, :amazon_q_oauth_application_id, if_exists: true + remove_column :ai_settings, :amazon_q_service_account_user_id, if_exists: true + remove_column :ai_settings, :amazon_q_ready, if_exists: true + remove_column :ai_settings, :amazon_q_role_arn, if_exists: true + end + end +end diff --git a/db/schema_migrations/20241120210519 b/db/schema_migrations/20241120210519 new file mode 100644 index 00000000000000..903304ef0e0e43 --- /dev/null +++ b/db/schema_migrations/20241120210519 @@ -0,0 +1 @@ +96e7602127a393034b48697a892e89830fc8fae140f9ace566e14f64ee603384 \ No newline at end of file diff --git a/db/structure.sql b/db/structure.sql index fd329190b34e01..7e415da43223ca 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -6248,7 +6248,12 @@ CREATE TABLE ai_settings ( singleton boolean DEFAULT true NOT NULL, created_at timestamp with time zone DEFAULT now() NOT NULL, updated_at timestamp with time zone DEFAULT now() NOT NULL, + amazon_q_oauth_application_id bigint, + amazon_q_service_account_user_id bigint, + amazon_q_ready boolean DEFAULT false NOT NULL, + amazon_q_role_arn text, CONSTRAINT check_3cf9826589 CHECK ((char_length(ai_gateway_url) <= 2048)), + CONSTRAINT check_a02bd8868c CHECK ((char_length(amazon_q_role_arn) <= 2048)), CONSTRAINT check_singleton CHECK ((singleton IS TRUE)) ); @@ -29305,6 +29310,10 @@ CREATE UNIQUE INDEX index_ai_feature_settings_on_feature ON ai_feature_settings CREATE UNIQUE INDEX index_ai_self_hosted_models_on_name ON ai_self_hosted_models USING btree (name); +CREATE INDEX index_ai_settings_on_amazon_q_oauth_application_id ON ai_settings USING btree (amazon_q_oauth_application_id); + +CREATE INDEX index_ai_settings_on_amazon_q_service_account_user_id ON ai_settings USING btree (amazon_q_service_account_user_id); + CREATE UNIQUE INDEX index_ai_settings_on_singleton ON ai_settings USING btree (singleton); CREATE INDEX index_ai_vectorizable_files_on_project_id ON ai_vectorizable_files USING btree (project_id); @@ -36020,6 +36029,9 @@ ALTER TABLE ONLY analytics_dashboards_pointers ALTER TABLE ONLY issues ADD CONSTRAINT fk_05f1e72feb FOREIGN KEY (author_id) REFERENCES users(id) ON DELETE SET NULL; +ALTER TABLE ONLY ai_settings + ADD CONSTRAINT fk_05f695565a FOREIGN KEY (amazon_q_oauth_application_id) REFERENCES oauth_applications(id) ON DELETE SET NULL; + ALTER TABLE ONLY merge_requests ADD CONSTRAINT fk_06067f5644 FOREIGN KEY (latest_merge_request_diff_id) REFERENCES merge_request_diffs(id) ON DELETE SET NULL; @@ -37370,6 +37382,9 @@ ALTER TABLE ONLY external_status_checks_protected_branches ALTER TABLE ONLY dast_profiles_pipelines ADD CONSTRAINT fk_cc206a8c13 FOREIGN KEY (dast_profile_id) REFERENCES dast_profiles(id) ON DELETE CASCADE; +ALTER TABLE ONLY ai_settings + ADD CONSTRAINT fk_cce81e0b9a FOREIGN KEY (amazon_q_service_account_user_id) REFERENCES users(id) ON DELETE SET NULL; + ALTER TABLE ONLY todos ADD CONSTRAINT fk_ccf0373936 FOREIGN KEY (author_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/ee/app/models/ai/setting.rb b/ee/app/models/ai/setting.rb index a9b8b449bddb38..277962dae7b20b 100644 --- a/ee/app/models/ai/setting.rb +++ b/ee/app/models/ai/setting.rb @@ -5,9 +5,22 @@ class Setting < ApplicationRecord self.table_name = "ai_settings" validates :ai_gateway_url, length: { maximum: 2048 }, allow_nil: true + validates :amazon_q_role_arn, length: { maximum: 2048 }, allow_nil: true + validate :validate_ai_gateway_url validate :validates_singleton + has_one :amazon_q_oauth_application, # rubocop:disable Rails/InverseOf -- This model is defined oustide of the project + class_name: 'Doorkeeper::Application', + foreign_key: 'id', + primary_key: 'amazon_q_oauth_application_id', + dependent: :nullify + has_one :amazon_q_service_account, # rubocop:disable Rails/InverseOf -- I don't think we need this? + class_name: 'User', + foreign_key: 'id', + primary_key: 'amazon_q_service_account_user_id', + dependent: :nullify + def self.instance first || create!(defaults) rescue ActiveRecord::RecordNotUnique, ActiveRecord::RecordInvalid diff --git a/ee/app/models/gitlab_subscriptions/features.rb b/ee/app/models/gitlab_subscriptions/features.rb index 917871ef585fba..93962e8618f6f7 100644 --- a/ee/app/models/gitlab_subscriptions/features.rb +++ b/ee/app/models/gitlab_subscriptions/features.rb @@ -186,15 +186,16 @@ class Features ].freeze ULTIMATE_FEATURES = %i[ + ai_agents ai_config_chat ai_features - ai_workflows - glab_ask_git_command - ai_agents ai_review_mr + ai_workflows + amazon_q api_discovery api_fuzzing auto_rollback + glab_ask_git_command cluster_receptive_agents cluster_image_scanning external_status_checks diff --git a/ee/config/feature_flags/wip/amazon_q_integration.yml b/ee/config/feature_flags/wip/amazon_q_integration.yml new file mode 100644 index 00000000000000..ccb47b547eae7c --- /dev/null +++ b/ee/config/feature_flags/wip/amazon_q_integration.yml @@ -0,0 +1,9 @@ +--- +name: amazon_q_integration +feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/508250 +introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/174949 +rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/508251 +milestone: '17.7' +group: group::ai framework +type: wip +default_enabled: false diff --git a/ee/lib/ai/amazon_q.rb b/ee/lib/ai/amazon_q.rb new file mode 100644 index 00000000000000..65c78eda346f59 --- /dev/null +++ b/ee/lib/ai/amazon_q.rb @@ -0,0 +1,18 @@ +# frozen_string_literal: true + +module Ai + module AmazonQ + # NOTE: This module is under development. See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/174614 + class << self + def feature_available? + ::Feature.enabled?(:amazon_q_integration, nil) && License.feature_available?(:amazon_q) + end + + def connected? + return false unless feature_available? + + Ai::Setting.instance.amazon_q_ready + end + end + end +end diff --git a/ee/spec/lib/ai/amazon_q_spec.rb b/ee/spec/lib/ai/amazon_q_spec.rb new file mode 100644 index 00000000000000..99478bc780b65d --- /dev/null +++ b/ee/spec/lib/ai/amazon_q_spec.rb @@ -0,0 +1,49 @@ +# frozen_string_literal: true + +require 'spec_helper' + +using RSpec::Parameterized::TableSyntax + +RSpec.describe Ai::AmazonQ, feature_category: :ai_abstraction_layer do + let(:application_settings) { ::Gitlab::CurrentSettings.current_application_settings } + + describe '#connected?' do + where(:q_available, :q_ready, :result) do + true | true | true + true | false | false + false | true | false + false | false | false + end + + with_them do + before do + Ai::Setting.instance.update!(amazon_q_ready: q_ready) + allow(described_class).to receive(:feature_available?).and_return(q_available) + end + + it 'returns the expected result' do + expect(described_class.connected?).to be result + end + end + end + + describe '#feature_available?' do + where(:feature_flag_enabled, :amazon_q_license_available, :result) do + true | true | true + true | false | false + false | true | false + false | false | false + end + + with_them do + before do + stub_licensed_features(amazon_q: amazon_q_license_available) + stub_feature_flags(amazon_q_integration: feature_flag_enabled) + end + + it 'returns the expected result' do + expect(described_class.feature_available?).to eq(result) + end + end + end +end diff --git a/ee/spec/models/ai/setting_spec.rb b/ee/spec/models/ai/setting_spec.rb index bf449c35cae9ab..26780c8487043d 100644 --- a/ee/spec/models/ai/setting_spec.rb +++ b/ee/spec/models/ai/setting_spec.rb @@ -3,6 +3,22 @@ require 'spec_helper' RSpec.describe Ai::Setting, feature_category: :ai_abstraction_layer do + describe 'associations' do + it 'has expected associations' do + is_expected.to have_one(:amazon_q_oauth_application) + .class_name('Doorkeeper::Application') + .with_foreign_key('id') + .with_primary_key('amazon_q_oauth_application_id') + .dependent(:nullify) + + is_expected.to have_one(:amazon_q_service_account) + .class_name('User') + .with_foreign_key('id') + .with_primary_key('amazon_q_service_account_user_id') + .dependent(:nullify) + end + end + describe 'validations' do subject(:setting) { described_class.instance } @@ -93,5 +109,7 @@ expect(described_class.count).to eq(1) end end + + it { is_expected.to validate_length_of(:amazon_q_role_arn).is_at_most(2048).allow_nil } end end -- GitLab From 7b85578e3fe7b01ed3e0591e665d8e5395f559cb Mon Sep 17 00:00:00 2001 From: Paul Slaughter Date: Fri, 6 Dec 2024 07:09:25 -0600 Subject: [PATCH 2/2] Polish Amazon Q db from review - Add if_not_exists to migration - Replace `has_one` with `belongs_to` - Add license feature to `GLOBAL_FEATURES` --- .../20241120210519_add_amazon_q_to_ai_settings.rb | 8 ++++---- ee/app/models/ai/setting.rb | 12 ++---------- ee/app/models/gitlab_subscriptions/features.rb | 1 + ee/spec/models/ai/setting_spec.rb | 13 ++----------- 4 files changed, 9 insertions(+), 25 deletions(-) diff --git a/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb b/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb index 1b78acf77b1814..5266be19bd8162 100644 --- a/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb +++ b/db/migrate/20241120210519_add_amazon_q_to_ai_settings.rb @@ -7,10 +7,10 @@ class AddAmazonQToAiSettings < Gitlab::Database::Migration[2.2] def up with_lock_retries do - add_column :ai_settings, :amazon_q_oauth_application_id, :bigint - add_column :ai_settings, :amazon_q_service_account_user_id, :bigint - add_column :ai_settings, :amazon_q_ready, :boolean, default: false, null: false - add_column :ai_settings, :amazon_q_role_arn, :text + add_column :ai_settings, :amazon_q_oauth_application_id, :bigint, if_not_exists: true + add_column :ai_settings, :amazon_q_service_account_user_id, :bigint, if_not_exists: true + add_column :ai_settings, :amazon_q_ready, :boolean, default: false, null: false, if_not_exists: true + add_column :ai_settings, :amazon_q_role_arn, :text, if_not_exists: true end add_concurrent_index :ai_settings, :amazon_q_oauth_application_id diff --git a/ee/app/models/ai/setting.rb b/ee/app/models/ai/setting.rb index 277962dae7b20b..c7434d38ece636 100644 --- a/ee/app/models/ai/setting.rb +++ b/ee/app/models/ai/setting.rb @@ -10,16 +10,8 @@ class Setting < ApplicationRecord validate :validate_ai_gateway_url validate :validates_singleton - has_one :amazon_q_oauth_application, # rubocop:disable Rails/InverseOf -- This model is defined oustide of the project - class_name: 'Doorkeeper::Application', - foreign_key: 'id', - primary_key: 'amazon_q_oauth_application_id', - dependent: :nullify - has_one :amazon_q_service_account, # rubocop:disable Rails/InverseOf -- I don't think we need this? - class_name: 'User', - foreign_key: 'id', - primary_key: 'amazon_q_service_account_user_id', - dependent: :nullify + belongs_to :amazon_q_oauth_application, class_name: 'Doorkeeper::Application', optional: true + belongs_to :amazon_q_service_account_user, class_name: 'User', optional: true def self.instance first || create!(defaults) diff --git a/ee/app/models/gitlab_subscriptions/features.rb b/ee/app/models/gitlab_subscriptions/features.rb index 93962e8618f6f7..705bdf61dd5dd5 100644 --- a/ee/app/models/gitlab_subscriptions/features.rb +++ b/ee/app/models/gitlab_subscriptions/features.rb @@ -11,6 +11,7 @@ class Features # For all other features, use `project.feature_available?` or `namespace.feature_available?` when possible. GLOBAL_FEATURES = %i[ admin_audit_log + amazon_q auditor_user custom_file_templates custom_project_templates diff --git a/ee/spec/models/ai/setting_spec.rb b/ee/spec/models/ai/setting_spec.rb index 26780c8487043d..8d6b2a56259ec4 100644 --- a/ee/spec/models/ai/setting_spec.rb +++ b/ee/spec/models/ai/setting_spec.rb @@ -5,17 +5,8 @@ RSpec.describe Ai::Setting, feature_category: :ai_abstraction_layer do describe 'associations' do it 'has expected associations' do - is_expected.to have_one(:amazon_q_oauth_application) - .class_name('Doorkeeper::Application') - .with_foreign_key('id') - .with_primary_key('amazon_q_oauth_application_id') - .dependent(:nullify) - - is_expected.to have_one(:amazon_q_service_account) - .class_name('User') - .with_foreign_key('id') - .with_primary_key('amazon_q_service_account_user_id') - .dependent(:nullify) + is_expected.to belong_to(:amazon_q_oauth_application).class_name('Doorkeeper::Application').optional + is_expected.to belong_to(:amazon_q_service_account_user).class_name('User').optional end end -- GitLab