embulk-input-mixpanel 0.5.15 → 0.6.1

data/lib/embulk/input/service/jql_service.rb

@@ -0,0 +1,276 @@
+require 'embulk/input/service/base_service'
+
+module Embulk
+  module Input
+    module Service
+      class JqlService < BaseService
+
+        FROM_DATE_PARAM = "params.from_date"
+        TO_DATE_PARAM = "params.to_date"
+
+        def validate_config
+          super
+
+          validate_jql_script
+          validate_fetch_days
+        end
+
+        def create_task
+          {
+            timezone: @config.param(:timezone, :string, default: ""),
+            api_secret: @config.param(:api_secret, :string),
+            jql_endpoint: endpoint,
+            dates: range,
+            incremental: @config.param(:incremental, :bool, default: true),
+            slice_range: @config.param(:slice_range, :integer, default: 7),
+            schema: @config.param(:columns, :array),
+            retry_initial_wait_sec: @config.param(:retry_initial_wait_sec, :integer, default: 1),
+            retry_limit: @config.param(:retry_limit, :integer, default: 5),
+            incremental_column: @config.param(:incremental_column, :string, default: nil),
+            latest_fetched_time: @config.param(:latest_fetched_time, :integer, default: 0),
+            jql_mode: true,
+            jql_script: @config.param(:jql_script, :string, nil)
+          }
+        end
+
+        def guess_columns
+          giveup_when_mixpanel_is_down
+          range = guess_range
+          Embulk.logger.info "Guessing schema using #{range.first}..#{range.last}"
+
+          client = create_client
+
+          sample_records = client.send_jql_script_small_dataset(parameters(@config.param(:jql_script, :string, nil), range.first, range.last))
+
+          validate_result(sample_records)
+
+          @incremental = @config.param(:incremental, :bool, default: true)
+          @incremental_column = @config.param(:incremental_column, :string, default: nil)
+          validate_result_contain_incremental_column(sample_records)
+
+          guess_from_records(sample_records)
+        end
+
+        def ingest(task, page_builder)
+          @dates = task[:dates]
+          @schema = task[:schema]
+          @timezone = task[:timezone]
+          @incremental_column = task[:incremental_column]
+          unless @incremental_column
+            Embulk.logger.warn "incremental_column should be specified when running in incremental mode to avoid duplicated"
+            Embulk.logger.warn "Use default value #{DEFAULT_TIME_COLUMN}"
+            @incremental_column = DEFAULT_TIME_COLUMN
+          end
+
+          @incremental = task[:incremental]
+          latest_fetched_time = task[:latest_fetched_time]
+
+          client = create_client
+
+          ignored_fetched_record_count = 0
+          next_fetched_time = latest_fetched_time
+          @dates.each_slice(task[:slice_range]) do |slice_dates|
+            Embulk.logger.info "Fetching date from #{slice_dates.first}..#{slice_dates.last}"
+            if preview?
+              records = client.send_jql_script_small_dataset(parameters(@config.param(:jql_script, :string, default: nil), slice_dates.first, slice_dates.last))
+            else
+              records = client.send_jql_script(parameters(task[:jql_script], slice_dates.first, slice_dates.last))
+            end
+            validate_result(records)
+            records.each do |record|
+              if @incremental
+                if @schema.map {|col| col["name"]}.include?(@incremental_column)
+                  record_incremental_column = record[@incremental_column]
+                  if record_incremental_column
+                    if record_incremental_column <= latest_fetched_time
+                      ignored_fetched_record_count += 1
+                      next
+                    else
+                      next_fetched_time = [record_incremental_column, next_fetched_time].max
+                    end
+                  end
+                else
+                  raise Embulk::ConfigError.new("Missing Incremental Field (<incremental_column>) in the returned dataset. Specify the correct Incremental Field value.")
+                end
+              end
+              values = extract_values(record)
+              page_builder.add(values)
+            end
+            break if preview?
+          end
+          Embulk.logger.info "Skip #{ignored_fetched_record_count} rows"
+          page_builder.finish
+
+          if task[:incremental] && !preview?
+            return create_task_report(next_fetched_time)
+          end
+          {}
+        end
+
+        def guess_range
+          time_zone = @config.param(:timezone, :string, default: "")
+          from_date = @config.param(:from_date, :string, default: default_guess_start_date(time_zone).to_s)
+          fetch_days = @config.param(:fetch_days, :integer, default: DEFAULT_FETCH_DAYS)
+
+          fetch_days = [fetch_days, DEFAULT_FETCH_DAYS].min
+
+          range = RangeGenerator.new(from_date, fetch_days, time_zone).generate_range
+          if range.empty?
+            return default_guess_start_date(time_zone)..(today(time_zone) - 1)
+          end
+          range
+        end
+
+        def guess_from_records(sample_props)
+          validate_result(sample_props)
+
+          begin
+            schema = Guess::SchemaGuess.from_hash_records(sample_props)
+            schema.map do |col|
+              result = {
+                name: col.name,
+                type: col.type,
+              }
+              if (col.name.eql? "time") || (col.eql? "last_seen")
+                result["format"] = col.format if col.format
+              end
+              result
+            end
+          rescue DataError
+            raise Embulk::ConfigError.new("Non-supported result #{sample_props}. Revise your JQL.")
+          end
+        end
+
+        def parameters(script, from_date, to_date)
+          {
+            params: params(from_date, to_date),
+            script: script
+          }
+        end
+
+        def adjust_timezone(epoch)
+          # Adjust timezone offset to get UTC time
+          # c.f. https://mixpanel.com/docs/api-documentation/exporting-raw-data-you-inserted-into-mixpanel#export
+          tz = TZInfo::Timezone.get(@timezone)
+
+          begin
+            if epoch.present?
+              offset = tz.period_for_local(epoch, true).offset.utc_total_offset
+              epoch - offset
+            end
+              # TZInfo::PeriodNotFound signals that there is no equivalent UTC time (for example,
+              # during the transition from standard time to daylight savings time when the clocks are moved forward and an hour is skipped).
+              # tz.local_time(2018, 3, 11, 2, 30, 0, 0)
+          rescue TZInfo::PeriodNotFound
+            epoch + 1.hour
+            offset = tz.period_for_local(epoch, true).offset.utc_total_offset
+            epoch - offset
+          end
+        end
+
+        def next_from_date(task_report)
+          next_to_date = Date.parse(task_report[:to_date])
+          {
+            from_date: next_to_date.to_s,
+            latest_fetched_time: task_report[:latest_fetched_time],
+          }
+        end
+
+        def endpoint
+          @config.param(:jql_endpoint, :string, default: Embulk::Input::MixpanelApi::Client::DEFAULT_JQL_ENDPOINT)
+        end
+
+        private
+
+        def create_task_report(next_fetched_time)
+          {
+            to_date: @dates.last || today(@timezone) - 1,
+            latest_fetched_time: next_fetched_time.to_s
+          }
+        end
+
+        def params(from_date, to_date)
+          {
+            from_date: from_date,
+            to_date: to_date
+          }
+        end
+
+        def extract_value(record, name)
+          case name
+          when NOT_PROPERTY_COLUMN
+            record[NOT_PROPERTY_COLUMN]
+          when "time"
+            if record["time"].present?
+              value = record["time"]
+              if value > 0
+                time = record["time"] / 1000
+                adjust_timezone(time)
+              else
+                value
+              end
+            end
+          when "last_seen"
+            if record["last_seen"].present?
+              value = record["last_seen"]
+              if value > 0
+                # last_seen format in ms
+                time = record["last_seen"] / 1000
+                adjust_timezone(time)
+              else
+                value
+              end
+            end
+          when @incremental_column
+            if record[@incremental_column].present?
+              value = record[@incremental_column]
+              if value > 0
+                # format in ms
+                time = record[@incremental_column] / 1000
+                adjust_timezone(time)
+              else
+                value
+              end
+            end
+          else
+            record[name]
+          end
+        end
+
+        def validate_result(records)
+          if records.is_a?(Array) && records.first.is_a?(Integer)
+            # incase using reduce, it only return the number of records
+            raise Embulk::ConfigError.new("Non-supported result. Revise your JQL.")
+          end
+        end
+
+        def validate_result_contain_incremental_column(records)
+          unless @incremental_column
+            Embulk.logger.warn "incremental_column should be specified when running in incremental mode to avoid duplicated"
+            Embulk.logger.warn "Use default value #{DEFAULT_TIME_COLUMN}"
+            @incremental_column = DEFAULT_TIME_COLUMN
+          end
+
+          if @incremental && records.length > 0 && !records[0].include?(@incremental_column)
+            raise Embulk::ConfigError.new("Missing Incremental Field (<incremental_column>) in the returned dataset. Specify the correct Incremental Field value.")
+          end
+        end
+
+        def validate_jql_script
+          jql_script = @config.param(:jql_script, :string, default: nil)
+          if jql_script.blank?
+            raise Embulk::ConfigError.new("JQL script shouldn't be empty or null")
+          end
+        end
+
+        def validate_fetch_days
+          fetch_days = @config.param(:fetch_days, :integer, default: nil)
+          if fetch_days && fetch_days <= 0
+            raise Embulk::ConfigError.new("fetch_days should be larger than 0")
+          end
+        end
+
+      end
+    end
+  end
+end

data/lib/timezone_validator.rb

@@ -9,7 +9,7 @@ class TimezoneValidator
       TZInfo::Timezone.get(@timezone)
     rescue => e
       Embulk.logger.error "'#{@timezone}' is invalid timezone"
-      raise Embulk::ConfigError.new e.message
+      raise Embulk::ConfigError.new ("Fail to identify timezone from '#{@timezone}':#{e.message}.")
     end
   end
 end

data/test/embulk/input/mixpanel_api/test_client.rb

@@ -9,22 +9,23 @@ module Embulk
         include OverrideAssertRaise
 
         API_SECRET = "api_secret".freeze
+        EXPORT_ENDPOINT = Embulk::Input::MixpanelApi::Client::DEFAULT_EXPORT_ENDPOINT
 
         def setup
-          @client = Client.new(API_SECRET)
+          @client = Client.new(API_SECRET, EXPORT_ENDPOINT)
           stub(Embulk).logger { ::Logger.new(IO::NULL) }
         end
 
         class TestKeepAlive < self
           def test_tcp_keepalive_enabled
-            client = Client.new(API_SECRET)
+            client = Client.new(API_SECRET, EXPORT_ENDPOINT)
             assert client.send(:httpclient).tcp_keepalive
           end
         end
 
         class TryToDatesTest < self
           def setup
-            @client = Client.new(API_SECRET)
+            @client = Client.new(API_SECRET, EXPORT_ENDPOINT)
           end
 
 
@@ -101,7 +102,6 @@ module Embulk
 
           def test_export_partial_with_error_json
             stub_client
-            # stub(@client).set_signatures(anything) {}
             stub_response(Struct.new(:code, :body).new(200, jsonl_dummy_responses+"\n{\"error\":"))
             records = []
             assert_raise MixpanelApi::IncompleteExportResponseError do
@@ -130,15 +130,6 @@ module Embulk
             end
           end
 
-          def test_retry_for_429_temporary_fail
-            stub_client
-            stub_response(failure_response(429))
-
-            assert_raise(RuntimeError) do
-              @client.export(params)
-            end
-          end
-
           class ExportSmallDataset < self
             def test_to_date_after_1_day
               to = (Date.parse(params["from_date"]) + 1).to_s
@@ -147,15 +138,6 @@ module Embulk
               @client.export_for_small_dataset(params)
             end
 
-            def test_retry_for_429_temporary_fail
-              stub_client
-              stub_response(failure_response(429))
-
-              assert_raise(RuntimeError) do
-                @client.export_for_small_dataset(params)
-              end
-            end
-
             def test_to_date_after_1_day_after_10_days_if_empty
               stub_client
               to1 = (Date.parse(params["from_date"]) + 1).to_s

data/test/embulk/input/{test_mixpanel.rb → service/test_export_service.rb}

@@ -1,12 +1,14 @@
 require "prepare_embulk"
 require "override_assert_raise"
 require "embulk/input/mixpanel"
+require "embulk/input/service/base_service"
+require "embulk/input/service/export_service"
 require "active_support/core_ext/time"
 require "json"
 
 module Embulk
   module Input
-    class MixpanelTest < Test::Unit::TestCase
+    class ExportServiceTest < Test::Unit::TestCase
       include OverrideAssertRaise
 
       API_SECRET = "api_secret".freeze
@@ -72,6 +74,7 @@ module Embulk
             type: "mixpanel",
             api_secret: API_SECRET,
             from_date: FROM_DATE,
+            timezone: TIMEZONE,
           }
 
           stub_export_all
@@ -90,7 +93,7 @@ module Embulk
           }
 
           stub_export_all
-          mock(Embulk.logger).info(/Guessing.*#{Regexp.escape Mixpanel.default_guess_start_date(TIMEZONE).to_s}/)
+          mock(Embulk.logger).info(/Guessing.*#{Regexp.escape Embulk::Input::Service::ExportService.new(config).default_guess_start_date(TIMEZONE).to_s}/)
 
           Mixpanel.guess(embulk_config(config))
         end
@@ -101,6 +104,7 @@ module Embulk
             type: "mixpanel",
             api_secret: API_SECRET,
             from_date: from_date,
+            timezone: TIMEZONE,
           }
 
           stub_export_all
@@ -113,11 +117,11 @@ module Embulk
           config = {
             type: "mixpanel",
             api_secret: API_SECRET,
-            timezone: TIMEZONE
+            timezone: TIMEZONE,
           }
 
           stub_export_all
-          mock(Embulk.logger).info(/Guessing.*#{Regexp.escape Mixpanel.default_guess_start_date(TIMEZONE).to_s}/)
+          mock(Embulk.logger).info(/Guessing.*#{Regexp.escape Embulk::Input::Service::ExportService.new(config).default_guess_start_date(TIMEZONE).to_s}/)
 
           Mixpanel.guess(embulk_config(config))
         end
@@ -126,7 +130,14 @@ module Embulk
           sample_records = records.map do |r|
             r.merge("properties" => {"time" => 1, "array" => [1, 2], "hash" => {foo: "FOO"}})
           end
-          actual = Mixpanel.guess_from_records(sample_records)
+
+          config = {
+            type: "mixpanel",
+            api_secret: API_SECRET,
+            timezone: TIMEZONE,
+          }
+
+          actual = Embulk::Input::Service::ExportService.new(config).guess_from_records(sample_records)
           assert actual.include?(name: "array", type: :json)
           assert actual.include?(name: "hash", type: :json)
         end
@@ -136,6 +147,7 @@ module Embulk
           config = {
             type: "mixpanel",
             api_secret: API_SECRET,
+            timezone: TIMEZONE,
           }
 
           assert_raise(Embulk::DataError) do
@@ -456,7 +468,7 @@ module Embulk
           where: 'properties["$os"] == "Windows"',
           bucket: "987",
         }
-        actual = Mixpanel.export_params(config)
+        actual = Embulk::Input::Service::ExportService.new(config).export_params
 
         assert_equal(expected, actual)
       end
@@ -534,7 +546,7 @@ module Embulk
             timezone: TIMEZONE,
             schema: schema,
             dates: DATES.to_a.map(&:to_s),
-            params: Mixpanel.export_params(embulk_config),
+            params: Mixpanel.service(embulk_config).export_params,
             fetch_unknown_columns: false,
             fetch_custom_properties: false,
             retry_initial_wait_sec: 0,
@@ -558,11 +570,16 @@ module Embulk
           super
           @page_builder = Object.new
           @plugin = Mixpanel.new(task, nil, nil, @page_builder)
-          stub(@plugin).fetch { records }
+          any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+            stub(klass).fetch { records }
+          end
+          # Embulk::Input::Service::ExportService.(:fetch => :records)
         end
 
         def test_preview
-          stub(@plugin).preview? { true }
+          any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+            stub(klass).preview? { true }
+          end
           mock(@page_builder).add(anything).times(records.length)
           mock(@page_builder).finish
 
@@ -570,7 +587,9 @@ module Embulk
         end
 
         def test_run
-          stub(@plugin).preview? { false }
+          any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+            stub(klass).preview? { false }
+          end
           mock(@page_builder).add(anything).times(records.length * 2)
           mock(@page_builder).finish
 
@@ -578,13 +597,16 @@ module Embulk
         end
 
         def test_timezone
-          stub(@plugin).preview? { false }
+          any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+            stub(klass).preview? { false }
+          end
           adjusted = record_epoch - timezone_offset_seconds
           mock(@page_builder).add(["FOO", adjusted, "event"]).times(records.length * 2)
           mock(@page_builder).finish
 
           @plugin.run
         end
+
         class PartialRunTest < self
           def setup_client
             any_instance_of(MixpanelApi::Client) do |klass|
@@ -598,7 +620,9 @@ module Embulk
 
           def test_run_with_allow_partial_false
             @plugin = Mixpanel.new(task.merge(allow_partial_import: false), nil, nil, @page_builder)
-            stub(@plugin).preview? {false}
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).preview? { false }
+            end
             assert_raise MixpanelApi::IncompleteExportResponseError do
               @plugin.run
             end
@@ -607,19 +631,25 @@ module Embulk
           def test_run_with_allow_partial_true
             @plugin = Mixpanel.new(task.merge(allow_partial_import: true), nil, nil, @page_builder)
             mock(@page_builder).finish
-            stub(@plugin).preview? {false}
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).preview? { false }
+            end
             @plugin.run
           end
         end
+
         class SliceRangeRunTest < self
 
           def test_default_slice_range
             plugin = Mixpanel.new(task.merge(slice_range: 2), nil, nil, @page_builder)
-            stub(plugin).preview? {false}
-            stub(plugin).fetch(["2015-02-22", "2015-02-23"],0){[]}
-            stub(plugin).fetch(["2015-02-24", "2015-02-25"],0){[]}
-            stub(plugin).fetch(["2015-02-26", "2015-02-27"],0){[]}
-            stub(plugin).fetch(["2015-02-28", "2015-03-01"],0){[]}
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).preview? { false }
+              stub(klass).fetch(["2015-02-22", "2015-02-23"],0,anything){[]}
+              stub(klass).fetch(["2015-02-24", "2015-02-25"],0,anything){[]}
+              stub(klass).fetch(["2015-02-26", "2015-02-27"],0,anything){[]}
+              stub(klass).fetch(["2015-02-28", "2015-03-01"],0,anything){[]}
+            end
+
             mock(@page_builder).finish
             plugin.run
           end
@@ -663,8 +693,10 @@ module Embulk
 
           def test_incremental_column_with_where
             page_builder = Object.new
-            plugin = Mixpanel.new(task.merge(params: task[:params].merge("where" => "abc==def"),latest_fetched_time: 1), nil, nil, page_builder)
-            stub(plugin).preview? {false}
+            plugin = Mixpanel.new(task.merge(params: task[:params].merge(where: "abc==def"),latest_fetched_time: 1), nil, nil, page_builder)
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).preview? { false }
+            end
             adjusted = record_epoch - timezone_offset_seconds
             mock(page_builder).add(["FOO", adjusted, "event"]).times(records.length * 2)
             mock(page_builder).finish
@@ -719,8 +751,10 @@ module Embulk
           def setup
             super
             @page_builder = Object.new
-            @plugin = Mixpanel.new(task, nil, nil, @page_builder)
-            stub(@plugin).fetch { [record] }
+            @plugin = Mixpanel.new(DataSource[task.to_a], nil, nil, @page_builder)
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).fetch { [record] }
+            end
           end
 
           def test_run
@@ -770,8 +804,10 @@ module Embulk
         class UnknownColumnsTest < self
           def setup
             @page_builder = Object.new
-            @plugin = Mixpanel.new(task, nil, nil, @page_builder)
-            stub(@plugin).fetch { records }
+            @plugin = Mixpanel.new(DataSource[task.to_a], nil, nil, @page_builder)
+            any_instance_of(Embulk::Input::Service::ExportService) do |klass|
+              stub(klass).fetch { records }
+            end
           end
 
           def test_run
@@ -840,7 +876,7 @@ module Embulk
           incremental_column: nil,
           schema: schema,
           dates: DATES.to_a.map(&:to_s),
-          params: Mixpanel.export_params(embulk_config),
+          params: Mixpanel.service(embulk_config).export_params,
           fetch_unknown_columns: false,
           fetch_custom_properties: false,
           retry_initial_wait_sec: 2,
@@ -880,6 +916,7 @@ module Embulk
           api_secret: API_SECRET,
           from_date: FROM_DATE,
           fetch_days: DAYS,
+          timezone: TIMEZONE,
           fetch_unknown_columns: false,
           fetch_custom_properties: false,
           retry_initial_wait_sec: 2,