Aggregation engines

class IssueAggregationEngine < Gitlab::Database::Aggregation::ActiveRecord::Engine
  filters do
    exact_match :project_id, :integer, description: 'Filter by project ID'
    exact_match :state, :string, description: 'Filter by issue state'
  end

  dimensions do
    column :author_id, :integer, description: 'Group by author'
    date_bucket :created_at, :datetime,
      parameters: { granularity: { in: %i[daily weekly monthly yearly], type: :string } },
      description: 'Group by creation date'
  end

  metrics do
    count description: 'Total number of issues'
    mean :weight, :float, description: 'Average issue weight'
  end
end

SELECT
  "issues"."author_id" AS aeq_author_id,
  date_trunc('month', "issues"."created_at") AS aeq_created_at,
  COUNT(*) AS aeq_total_count,
  AVG("issues"."weight") AS aeq_mean_weight
FROM "issues"
WHERE "issues"."project_id" IN (1, 2, 3)
  AND "issues"."state" IN ('opened')
GROUP BY aeq_author_id, aeq_created_at
ORDER BY aeq_author_id, aeq_created_at

class SessionAnalyticsEngine < Gitlab::Database::Aggregation::ClickHouse::Engine
  self.table_name = 'sessions'

  filters do
    exact_match :flow_type, :string, description: 'Filter by flow type'
    range :created_at, :datetime, description: 'Filter by creation date'
  end

  dimensions do
    column :flow_type, :string, description: 'Group by flow type'
    date_bucket :created_at, :datetime,
      parameters: { granularity: { in: %i[daily weekly monthly], type: :string } },
      description: 'Group by date'
  end

  metrics do
    count description: 'Total sessions'
    count :completed, :integer,
      expression: -> { Arel.sql('1') },
      if: -> { Arel.sql('finished_at IS NOT NULL') },
      description: 'Completed sessions'
    mean :duration, :float,
      expression: -> { Arel.sql('finished_at - created_at') },
      if: -> { Arel.sql('finished_at IS NOT NULL') },
      description: 'Average session duration'
    rate :completion,
      numerator_if: -> { Arel.sql('finished_at IS NOT NULL') },
      description: 'Session completion rate'
    quantile :duration, :float,
      expression: -> { Arel.sql('finished_at - created_at') },
      parameters: { quantile: { type: :float, description: 'Quantile value (0.0-1.0)' } },
      description: 'Duration percentile'
  end
end

-- metacode query to emphasize on query structure
SELECT dimensions, metrics
FROM (
  SELECT
    primary_key_columns,
    dimensions_expressions,
    metrics_expressions,
  FROM source_table
  WHERE filters
  GROUP BY ALL
) ch_aggregation_inner_query
GROUP BY ALL
ORDER BY orders

SELECT
  `ch_aggregation_inner_query`.`aeq_flow_type` AS aeq_flow_type,
  toStartOfInterval(
    `ch_aggregation_inner_query`.`aeq_created_at`,
    INTERVAL 1 month
  ) AS aeq_created_at,
  COUNT(*) AS aeq_total_count,
  countIf(`ch_aggregation_inner_query`.`aeq_completed_secondary` = 1) AS aeq_completed_count,
  avgIf(
    `ch_aggregation_inner_query`.`aeq_mean_duration`,
    `ch_aggregation_inner_query`.`aeq_mean_duration_secondary` = 1
  ) AS aeq_mean_duration,
  countIf(`ch_aggregation_inner_query`.`aeq_completion_rate` = 1) / COUNT(*) AS aeq_completion_rate,
  quantile(0.5)(`ch_aggregation_inner_query`.`aeq_duration_quantile`) AS aeq_duration_quantile
FROM (
  SELECT
    `sessions`.`flow_type` AS aeq_flow_type,
    `sessions`.`created_at` AS aeq_created_at,
    finished_at IS NOT NULL AS aeq_completed_secondary,
    finished_at - created_at AS aeq_mean_duration,
    finished_at IS NOT NULL AS aeq_mean_duration_secondary,
    finished_at IS NOT NULL AS aeq_completion_rate,
    finished_at - created_at AS aeq_duration_quantile,
    `sessions`.`user_id`,
    `sessions`.`session_id`
  FROM `sessions`
  WHERE `sessions`.`created_at` BETWEEN '2024-01-01' AND '2024-12-31'
  GROUP BY ALL
) ch_aggregation_inner_query
GROUP BY ALL
ORDER BY aeq_flow_type, aeq_created_at

metrics do
  retained_count :returning_users, :integer, -> { sql('user_id') }, over: :timestamp,
    description: 'Users present in both the current and previous period'
end

metrics do
  lagged_count :previous_period_users, :integer, -> { sql('user_id') }, over: :timestamp,
    description: 'Distinct users in the previous period'
end

filters do
  metric_exact_match :total_count, :integer
end

Gitlab::Database::Aggregation::Request.new(
  filters: [{ identifier: :total_count, values: [1, 2] }],
  dimensions: [{ identifier: :user_id }],
  metrics: [{ identifier: :total_count }]
)

filters do
  metric_range :total_count, :integer
  metric_range :duration_quantile, :float
end

Gitlab::Database::Aggregation::Request.new(
  filters: [
    { identifier: :duration_quantile, parameters: { quantile: 0.1 }, values: 200..nil }
  ],
  dimensions: [{ identifier: :user_id }],
  metrics: [{ identifier: :duration_quantile, parameters: { quantile: 0.1 } }]
)

transient(:duration) do
  sql("dateDiff('seconds', anyIfMerge(created_event_at), anyIfMerge(finished_event_at))")
end

transient(:is_finished) { sql('anyIfMerge(finished_event_at) IS NOT NULL') }

metrics do
  mean :duration, :float, transient(:duration),
    description: 'Average session duration in seconds'

  count :finished, if: transient(:is_finished),
    description: 'Number of finished sessions'
end

request = Gitlab::Database::Aggregation::Request.new(
  filters: [
    { identifier: :project_id, values: [1, 2, 3] },
    { identifier: :state, values: ['opened'] }
  ],
  dimensions: [
    { identifier: :author_id },
    { identifier: :created_at, parameters: { granularity: 'monthly' } },
    { identifier: :created_at, parameters: { granularity: 'weekly' } },
  ],
  metrics: [
    { identifier: :total_count },
    { identifier: :mean_weight }
  ],
  order: [
    { identifier: :total_count, direction: :desc } # order identifier must reference dimension or metric.
  ]
)

engine = IssueAggregationEngine.new(context: { scope: Issue.all })
response = engine.execute(request)

if response.success?
  puts "Success: #{response.payload[:data].to_a.inspect}"
else
  puts "Errors: #{response.errors}"
end

┌────────────────────────────────────────────────────┐
│                        Request                     │
│  (metrics, dimensions, filters, order)             │
└────────────────────────────────────────────────────┘
                            │
                            ▼
┌────────────────────────────────────────────────────┐
│                        QueryPlan                   │
│  (validates request, builds plan parts)            │
└────────────────────────────────────────────────────┘
                            │
                            ▼
┌────────────────────────────────────────────────────┐
│                         Engine                     │
│  (executes query plan, returns AggregationResult)  │
└────────────────────────────────────────────────────┘
                            │
                            ▼
┌────────────────────────────────────────────────────┐
│                   AggregationResult                │
│ implements Enumerable to access formatted results  │
└────────────────────────────────────────────────────┘

module Types
  class ProjectType < BaseObject
    extend Gitlab::Database::Aggregation::Graphql::Mounter

    mount_aggregation_engine(
      IssueAggregationEngine,
      field_name: 'issue_analytics',
      description: 'Issue analytics aggregation'
    ) do
      # Define base aggregation scope. Build your own scope or inherit one from parent object.
      def aggregation_scope
        object.issues
      end
    end
  end
end

mount_aggregation_engine(
  IssueAggregationEngine,
  field_name: 'issue_analytics',
  description: 'Issue analytics aggregation',
  authorize: :read_project
) do
  # authorize :read_project - this also supported.
  def aggregation_scope
    object.issues
  end
end

query IssueAnalytics($projectId: ID!) {
  project(fullPath: $projectId) {
    issueAnalytics(
      state: ["opened", "closed"]
      createdAtFrom: "2024-01-01"
      createdAtTo: "2024-12-31"
    ) {
      aggregated(
        totalCountFrom: 5
        orderBy: [{ identifier: "totalCount", direction: DESC }]
        first: 10
      ) {
        nodes {
          dimensions {
            createdAt(granularity: "monthly")
          }
          totalCount
          meanWeight
          highQuantile: durationQuantile(0.9)
          medianQuantile: durationQuantile(0.5)
        }
        pageInfo {
          hasNextPage
          endCursor
        }
      }
    }
  }
}

module Types
  class ProjectType < BaseObject
    extend Gitlab::Database::Aggregation::Graphql::Mounter

    mount_aggregation_engine(IssueAggregationEngine) do
      # Other configuration options...
      # Custom validation logic
      def validate_request!(engine_request)
        if engine_request.dimensions.empty?
          raise GraphQL::ExecutionError, 'At least one dimension must be specified'
        end
      end
    end
  end
end

class AgentPlatformSessions < Gitlab::Database::Aggregation::ClickHouse::Engine
  dimensions do
    column :flow_type, :string, description: 'Type of session'
    column :user_id, :integer, description: 'Session owner', association: true
  end
end

dimensions do
  column :author_id, :integer,
    description: 'Issue author',
    association: { model: User }
    # or model and GraphQL type
    # association: { model: User, graphql_type: Types::CurrentUserType }
end

query {
  project(fullPath: "gitlab-org/gitlab") {
    aiUsage {
      agentPlatformSessions {
        aggregated {
          nodes {
            dimensions {
              userId  # Returns: 123 (integer)
            }
          }
        }
      }
    }
  }
}

query {
  project(fullPath: "gitlab-org/gitlab") {
    aiUsage {
      agentPlatformSessions(
        userId: [1, 2]  # Filter still uses original dimension identifier
      ) {
        aggregated(
          orderBy: [{ identifier: "user", direction: DESC }]  # Order uses association name
        ) {
          nodes {
            dimensions {
              user {  # Returns: full User object
                id
                username
                name
              }
            }
          }
        }
      }
    }
  }
}