Make UserCohortsService more understandable

1. Extract out into several methods.
2. Add more comments describing the data and the shape of the data.

2 files changed, 73 insertions, 39 deletions

diff --git a/app/services/user_cohorts_service.rb b/app/services/user_cohorts_service.rb
index 7f84b6a0634..6545ceffec6 100644
--- a/app/services/user_cohorts_service.rb
+++ b/app/services/user_cohorts_service.rb
@@ -1,41 +1,32 @@
 class UserCohortsService
-  def initialize
-  end
+  MONTHS_INCLUDED = 12
 
-  def execute(months_included)
-    if Gitlab::Database.postgresql?
-      created_at_month = "CAST(DATE_TRUNC('month', created_at) AS date)"
-      current_sign_in_at_month = "CAST(DATE_TRUNC('month', current_sign_in_at) AS date)"
-    elsif Gitlab::Database.mysql?
-      created_at_month = "STR_TO_DATE(DATE_FORMAT(created_at, '%Y-%m-01'), '%Y-%m-%d')"
-      current_sign_in_at_month = "STR_TO_DATE(DATE_FORMAT(current_sign_in_at, '%Y-%m-01'), '%Y-%m-%d')"
-    end
+  # Get a hash that looks like:
+  #
+  #     {
+  #       month => {
+  #         months: [3, 2, 1],
+  #         total: 3
+  #         inactive: 0
+  #      },
+  #      etc.
+  #
+  # The `months` array is always from oldest to newest, so it's always
+  # non-strictly decreasing from left to right.
+  #
+  def execute
+    cohorts = {}
+    months = Array.new(MONTHS_INCLUDED) { |i| i.months.ago.beginning_of_month.to_date }
 
-    counts_by_month =
-      User
-        .where('created_at > ?', months_included.months.ago.end_of_month)
-        .group(created_at_month, current_sign_in_at_month)
-        .reorder("#{created_at_month} ASC", "#{current_sign_in_at_month} DESC")
-        .count
+    MONTHS_INCLUDED.times do
+      created_at_month = months.last
+      activity_months = running_totals(months, created_at_month)
 
-    cohorts = {}
-    months = Array.new(months_included) { |i| i.months.ago.beginning_of_month.to_date }
-
-    months_included.times do
-      month = months.last
-      inactive = counts_by_month[[month, nil]] || 0
-
-      # Calculate a running sum of active users, so users active in later months
-      # count as active in this month, too. Start with the most recent month
-      # first, for calculating the running totals, and then reverse for
-      # displaying in the table.
-      activity_months =
-        months
-          .map { |activity_month| counts_by_month[[month, activity_month]] }
-          .reduce([]) { |result, total| result << result.last.to_i + total.to_i }
-          .reverse
-
-      cohorts[month] = {
+      # Even if no users registered in this month, we always want to have a
+      # value to fill in the table.
+      inactive = counts_by_month[[created_at_month, nil]].to_i
+
+      cohorts[created_at_month] = {
         months: activity_months,
         total: activity_months.first,
         inactive: inactive
@@ -46,4 +37,51 @@ class UserCohortsService
 
     cohorts
   end
+
+  private
+
+  # Calculate a running sum of active users, so users active in later months
+  # count as active in this month, too. Start with the most recent month first,
+  # for calculating the running totals, and then reverse for displaying in the
+  # table.
+  def running_totals(all_months, created_at_month)
+    all_months
+      .map { |activity_month| counts_by_month[[created_at_month, activity_month]] }
+      .reduce([]) { |result, total| result << result.last.to_i + total.to_i }
+      .reverse
+  end
+
+  # Get a hash that looks like:
+  #
+  #     {
+  #       [created_at_month, current_sign_in_at_month] => count,
+  #       [created_at_month, current_sign_in_at_month_2] => count_2,
+  #       # etc.
+  #     }
+  #
+  # created_at_month can never be nil, but current_sign_in_at_month can (when a
+  # user has never logged in, just been created). This covers the last twelve
+  # months.
+  #
+  def counts_by_month
+    @counts_by_month ||=
+      begin
+        created_at_month = column_to_date('created_at')
+        current_sign_in_at_month = column_to_date('current_sign_in_at_month')
+
+        User
+          .where('created_at > ?', MONTHS_INCLUDED.months.ago.end_of_month)
+          .group(created_at_month, current_sign_in_at_month)
+          .reorder("#{created_at_month} ASC", "#{current_sign_in_at_month} ASC")
+          .count
+      end
+  end
+
+  def column_to_date(column)
+    if Gitlab::Database.postgresql?
+      "CAST(DATE_TRUNC('month', #{column}) AS date)"
+    elsif Gitlab::Database.mysql?
+      "STR_TO_DATE(DATE_FORMAT(#{column}, '%Y-%m-01'), '%Y-%m-%d')"
+    end
+  end
 end
diff --git a/spec/services/user_cohorts_service_spec.rb b/spec/services/user_cohorts_service_spec.rb
index 8d8d0de31cd..4db8d577d79 100644
--- a/spec/services/user_cohorts_service_spec.rb
+++ b/spec/services/user_cohorts_service_spec.rb
@@ -32,11 +32,7 @@ describe UserCohortsService do
         month_start(0) => { months: [2], total: 2, inactive: 1 }
       }
 
-      result = described_class.new.execute(12)
-
-      expect(result.length).to eq(12)
-      expect(result.keys).to all(be_a(Date))
-      expect(result).to eq(expected)
+      expect(described_class.new.execute).to eq(expected)
     end
   end
 end