dwh 0.1.0

data/lib/dwh/adapters/druid.rb

@@ -0,0 +1,267 @@
+module DWH
+  module Adapters
+    # Druid adapter.
+    #
+    # Generally, adapters should be created using {DWH::Factory#create DWH.create}. Where a configuration
+    # is passed in as options hash or argument list.
+    #
+    # @example Basic connection with required only options
+    #   DWH.create(:druid, {host: 'localhost',port: 8080, protocol: 'http'})
+    #
+    # @example Connect with SSL and basic authorization
+    #   DWH.create(:druid, {host: 'localhost',port: 8080, protocol: 'http',
+    #       basic_auth: 'BASE_64 encoded authorization key'
+    #   })
+    #
+    # @example Sending custom client name and user information
+    #   DWH.create(:druid, {host: 'localhost',port: 8080,
+    #     client_name: 'Strata CLI', extra_connection_params: {
+    #       context: {
+    #         user: 'Ajo',
+    #         team: 'Engineering'
+    #       }
+    #     }})
+    class Druid < Adapter
+      DRUID_STATUS = '/status'.freeze
+      DRUID_DATASOURCES = '/druid/coordinator/v1/datasources'.freeze
+      DRUID_SQL = '/druid/v2/sql/'.freeze
+      COLUMNS_FOR_TABLE = '"COLUMN_NAME","DATA_TYPE", "NUMERIC_PRECISION", "NUMERIC_SCALE", "CHARACTER_MAXIMUM_LENGTH"'.freeze
+
+      config :protocol, String, required: true, default: 'http', message: 'must be http or https', allowd: %w[http https]
+      config :host, String, required: true, message: 'server host ip address or domain name'
+      config :port, Integer, required: true, default: 8081, message: 'port to connect to'
+      config :query_timeout, Integer, required: false, default: 600, message: 'query execution timeout in seconds'
+      config :open_timeout, Integer, required: false, default: nil, message: 'how long to wait to connect'
+      config :client_name, String, default: 'DWH Ruby Gem', message: 'client_name will be passed in the context object'
+      config :basic_auth, String, required: false, message: 'authorization key sent in the header'
+
+      # (see Adapter#connection)
+      def connection
+        return @connection if @connection
+
+        @connection = Faraday.new(
+          url: "#{config[:protocol]}://#{config[:host]}:#{config[:port]}",
+          headers: {
+            'Content-Type' => 'application/json',
+            **(config[:basic_auth] ? { 'Authorization' => "Basic #{config[:basic_auth]}" } : {})
+          },
+          request: {
+            timeout: config[:query_timeout],
+            open_timeout: config[:open_timeout],
+            context: {
+              client_name: config[:client_name]
+            }
+          }.merge(extra_connection_params)
+        )
+
+        @connection
+      end
+
+      # (see Adapter#test_connection)
+      def test_connection(raise_exception: false)
+        res = connection.get(DRUID_STATUS)
+        unless res.success?
+          raise ConnectionError, res.body if raise_exception
+
+          false
+        end
+
+        true
+      rescue Faraday::ConnectionFailed => e
+        raise ConnectionError, e.message if raise_exception
+
+        false
+      end
+
+      # (see Adapter#tables)
+      def tables
+        resp = connection.get(DRUID_DATASOURCES) do |req|
+          req.options.timeout = 30
+        end
+        JSON.parse resp.body
+      end
+
+      # Date column will default to __time. If the datasource,
+      # does not have a date column please set it to nil
+      # @param table [String] table name
+      # @param date_column [String] optional date column
+      # @see Adapter#stats
+      def stats(table, date_column: '__time')
+        sql = <<-SQL
+        SELECT
+        count(*) ROW_COUNT
+        #{date_column.nil? ? nil : ", min(#{date_column}) DATE_START"}
+        #{date_column.nil? ? nil : ", max(#{date_column}) DATE_END"}
+        FROM "#{table}"
+        SQL
+
+        result = execute(sql)
+
+        TableStats.new(
+          row_count: result[0][0],
+          date_start: result[0][1],
+          date_end: result[0][2]
+        )
+      end
+
+      # Marks unused segments of a datasource/table as unused
+      # @param table [String] datasource/table name
+      # @param interval [String] date interval in the format of from_date/to_date
+      #   as valid ISO timestamps
+      def drop_unused_segments(table, interval)
+        url = "/druid/coordinator/v1/datasources/#{table}/markUnused"
+
+        logger.debug '=== Dropping Segments ==='
+
+        response = connection.post(url) do |req|
+          req.headers['Content-Type'] = 'application/json'
+          req.body = { interval: interval }.to_json
+        end
+
+        logger.debug response.status
+      end
+
+      # (see Adapter#metadata)
+      def metadata(table)
+        sql = <<-SQL
+        SELECT #{COLUMNS_FOR_TABLE} FROM INFORMATION_SCHEMA.COLUMNS
+        WHERE TABLE_SCHEMA = 'druid' AND TABLE_NAME = '#{table}'
+        SQL
+
+        stats = stats(table)
+        db_table = Table.new 'table', table_stats: stats
+        cols = execute(sql, format: :object)
+        st = table_druid_schema_types(table, stats.date_end)
+
+        cols.each do |col|
+          db_table << Column.new(
+            name: col['COLUMN_NAME'],
+            schema_type: st[:metrics].include?(col['COLUMN_NAME']) ? 'measure' : 'dimension',
+            data_type: col['DATA_TYPE'],
+            precision: col['NUMERIC_PRECISION'],
+            scale: col['NUMERIC_SCALE'],
+            max_char_length: col['CHARACTER_MAXIMUM_LENGTH']
+          )
+        end
+
+        db_table
+      end
+
+      # (see Adapter#execute)
+      def execute(sql, format: :array, retries: 0)
+        format = format.to_sym
+        result_format = format == :native ? 'array' : format.to_s
+        resp = with_debug(sql) do
+          with_retry(retries) do
+            connection.post(DRUID_SQL) do |req|
+              req.headers['Content-Type'] = 'application/json'
+              req.body = {
+                query: sql,
+                resultFormat: result_format,
+                context: { sqlTimeZone: 'Etc/UTC' }
+              }.merge(extra_query_params)
+                         .to_json
+            end
+          end
+        end
+
+        raise ExecutionError, "Could not execute #{sql}: \n #{resp.body}" if resp.status != 200
+
+        if format == :native
+          resp
+        else
+          format == :csv ? resp.body : JSON.parse(resp.body)
+        end
+      end
+
+      # (see Adapter#execute_stream)
+      def execute_stream(sql, io, stats: nil, retries: 0)
+        resp = with_debug(sql) do
+          with_retry(retries) do
+            connection.post(DRUID_SQL) do |req|
+              req.headers['Content-Type'] = 'application/json'
+              req.body = {
+                query: sql,
+                resultFormat: 'csv',
+                header: true
+                # added timezone here due to druid bug
+                # where date sub query joins failed without it.
+                # context: { sqlTimeZone: 'Etc/UTC'}
+              }.merge(extra_query_params).to_json
+
+              parseable_row = ''
+              req.options.on_data = proc do |chunk, _|
+                handle_streaming_chunk(io, chunk, stats, parseable_row)
+              end
+            end
+          end
+        end
+
+        io.rewind
+        # Raise exception on failed runs
+        raise ExecutionError, io.read unless resp.success?
+
+        io
+      end
+
+      # (see Adapter#stream)
+      def stream(sql, &block)
+        on_data_calls = 0
+        with_debug(sql) do
+          connection.post(DRUID_SQL) do |req|
+            req.headers['Content-Type'] = 'application/json'
+            req.body = { query: sql, resultFormat: 'csv' }.to_json
+            req.options.on_data = proc do |chunk, _chunk_size|
+              block.call chunk.force_encoding('utf-8')
+              on_data_calls += 1
+            end
+          end
+        end
+
+        on_data_calls
+      end
+
+      protected
+
+      def table_druid_schema_types(table, last_interval_start_date)
+        end_date = last_interval_start_date + 1
+        start_date = last_interval_start_date
+        url_friendly_interval = "#{start_date.strftime('%Y-%m-%d')}_#{end_date.strftime('%Y-%m-%d')}"
+        url = "/druid/coordinator/v1/datasources/#{table}/intervals/#{url_friendly_interval}?full"
+
+        resp = connection.get(url) do |req|
+          req.options.timeout = 30
+        end
+
+        raise ExecutionError, "Could not fetch druid schema types: \n #{resp.body}" if resp.status != 200
+
+        res = JSON.parse(resp.body)
+        meta = res.flatten[1].flatten(4)[1]['metadata']
+        {
+          dimensions: meta['dimensions'].split(','),
+          metrics: meta['metrics'].split(',')
+        }
+      end
+
+      def handle_streaming_chunk(io, chunk, stats, parseable_row)
+        io.write chunk.rstrip.force_encoding('utf-8')
+
+        parseable_row += chunk
+        process_streaming_rows(parseable_row, chunk, stats)
+      end
+
+      def process_streaming_rows(parseable_row, chunk, stats)
+        return if stats.nil? || stats&.limit_reached?
+
+        rows = CSV.parse(parseable_row, skip_blanks: true)
+        rows.each_with_index do |row, index|
+          # skip header rows in stats collector
+          stats << row unless index.zero? && stats.total_rows.zero?
+        end
+        parseable_row.clear
+      rescue CSV::MalformedCSVError
+        logger.debug("Unparseable:\n #{chunk}")
+      end
+    end
+  end
+end

data/lib/dwh/adapters/duck_db.rb

@@ -0,0 +1,235 @@
+module DWH
+  module Adapters
+    # DuckDb adapter.
+    #
+    # This requires the ruby {https://github.com/suketa/ruby-duckdb DuckDb} gem.  Installation
+    # is a bit complex. Please follow the guide on the gems page to make sure
+    # you have DuckDb installed as required before installing the gem.
+    #
+    # Generally, adapters should be created using {DWH::Factory#create DWH.create}. Where a configuration
+    # is passed in as options hash or argument list.
+    #
+    # @example Basic connection with required only options
+    #   DWH.create(:duckdb, {file: 'path/to/my/duckdb' })
+    #
+    # @example Open in read only mode. {https://duckdb.org/docs/stable/configuration/overview#configuration-reference config docs}
+    #   DWH.create(:duckdb, {file: 'path/to/my/duckdb' ,duck_config: { access_mode: "READ_ONLY"}})
+    class DuckDb < Adapter
+      config :file, String, required: true, message: 'path/to/duckdb/db'
+      config :schema, String, required: false, default: 'main', message: 'schema defaults to main'
+      config :duck_config, Hash, required: false, message: 'hash of valid DuckDb configuration options'
+
+      # (see Adapter#connection)
+      def connection
+        return @connection if @connection
+
+        if self.class.databases.key?(config[:file])
+          @db = self.class.databases[config[:file]]
+        else
+          ducked_config = DuckDB::Config.new
+          if config.key?(:duck_config)
+            config[:duck_config].each do |key, val|
+              ducked_config[key.to_s] = val
+            end
+          end
+          @db = DuckDB::Database.open(config[:file], ducked_config)
+          self.class.databases[config[:file]] = @db
+        end
+
+        @connection = @db.connect
+
+        @connection
+      rescue StandardError => e
+        raise ConfigError, e.message
+      end
+
+      def self.databases
+        @databases ||= {}
+      end
+
+      def self.open_databases
+        databases.size
+      end
+
+      # DuckDB is an in process database so we don't want to
+      # open multiple instances of the same db in memory. Rather,
+      # we open one instance but many connections. Use this
+      # method to close them all.
+      def self.close_all
+        databases.each do |key, db|
+          db.close
+          databases.delete(key)
+        end
+      end
+
+      # This disconnects the current connection but
+      # the db is still in process and can be reconnected
+      # to.
+      #
+      # (see Adapter#close)
+      def close
+        connection.disconnect
+        @connection = nil
+      end
+
+      # (see Adapter#test_connection)
+      def test_connection(raise_exception: false)
+        connection
+        true
+      rescue StandardError => e
+        raise ConnectionError, e.message if raise_exception
+
+        false
+      end
+
+      # (see Adapter#tables)
+      def tables(**qualifiers)
+        catalog, schema = qualifiers.values_at(:catalog, :schema)
+        sql = 'SELECT table_name FROM duckdb_tables'
+
+        where = []
+        where << "database_name = '#{catalog}'" if catalog
+
+        where << if schema
+                   "schema_name = '#{schema}'"
+                 else
+                   "schema_name = '#{config[:schema]}'"
+                 end
+
+        res = execute("#{sql} WHERE #{where.join(' AND ')}")
+        res.flatten
+      end
+
+      # (see Adapter#stats)
+      def stats(table, date_column: nil, **qualifiers)
+        qualifiers[:schema] = config[:schema] unless qualifiers[:schema]
+        db_table = Table.new table, **qualifiers
+
+        sql = <<-SQL
+        SELECT count(*) ROW_COUNT
+        #{date_column.nil? ? nil : ", min(#{date_column}) DATE_START"}
+        #{date_column.nil? ? nil : ", max(#{date_column}) DATE_END"}
+        FROM #{db_table.fully_qualified_table_name}
+        SQL
+
+        result = execute(sql)
+        TableStats.new(
+          row_count: result.first[0],
+          date_start: result.first[1],
+          date_end: result.first[2]
+        )
+      end
+
+      # (see Adapter#metadata)
+      def metadata(table, **qualifiers)
+        db_table = Table.new table, **qualifiers
+        sql = 'SELECT column_name, data_type, character_maximum_length, numeric_precision,numeric_scale FROM duckdb_columns'
+
+        where = ["table_name = '#{db_table.physical_name}'"]
+        where << "database_name = '#{db_table.catalog}'" if db_table.catalog
+
+        where << if db_table.schema
+                   "schema_name = '#{db_table.schema}'"
+                 else
+                   "schema_name = '#{config[:schema]}'"
+                 end
+
+        cols = execute("#{sql} WHERE #{where.join(' AND ')}")
+        cols.each do |col|
+          db_table << Column.new(
+            name: col[0],
+            data_type: col[1],
+            precision: col[3],
+            scale: col[4],
+            max_char_length: col[2]
+          )
+        end
+
+        db_table
+      end
+
+      # True if the configuration was setup with a schema.
+      def schema?
+        config[:schema].present?
+      end
+
+      # (see Adapter#execute)
+      def execute(sql, format: :array, retries: 0)
+        begin
+          result = with_debug(sql) { with_retry(retries) { connection.query(sql) } }
+        rescue StandardError => e
+          raise ExecutionError, e.message
+        end
+
+        format = format.downcase if format.is_a?(String)
+        case format.to_sym
+        when :array
+          result.to_a
+        when :object
+          result_to_hash(result)
+        when :csv
+          result_to_csv(result)
+        when :native
+          result
+        else
+          raise UnsupportedCapability, "Unsupported format: #{format} for this #{name}"
+        end
+      end
+
+      # (see Adapter#execute_stream)
+      def execute_stream(sql, io, stats: nil, retries: 0)
+        with_debug(sql) do
+          with_retry(retries) do
+            result = connection.query(sql)
+            io.write(CSV.generate_line(result.columns.map(&:name)))
+            result.each do |row|
+              stats << row unless stats.nil?
+              io.write(CSV.generate_line(row))
+            end
+          end
+        end
+
+        io.rewind
+        io
+      rescue StandardError => e
+        raise ExecutionError, e.message
+      end
+
+      # (see Adapter#stream)
+      def stream(sql, &block)
+        with_debug(sql) do
+          result = connection.query(sql)
+          result.each do |row|
+            block.call(row)
+          end
+        end
+      end
+
+      def valid_config?
+        super
+        require 'duckdb'
+      rescue LoadError
+        raise ConfigError, "Required 'duckdb' gem missing. Please add it to your Gemfile."
+      end
+
+      private
+
+      def result_to_hash(result)
+        columns = result.columns.map(&:name)
+
+        result.each.map do |row|
+          columns.zip(row).to_h
+        end
+      end
+
+      def result_to_csv(result)
+        CSV.generate do |csv|
+          csv << result.columns.map(&:name)
+          result.each do |row|
+            csv << row
+          end
+        end
+      end
+    end
+  end
+end