s3_selector 0.1.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 60db797c466d34b7363115daa3e6baecfdf5ab12e5362d6fc13c6fd2a2fcc977
+  data.tar.gz: d3e600200aeb2087b6186bea76445415fcff1d8d444fc7ba9048434032810eb1
+SHA512:
+  metadata.gz: 530d505fb1c38cd4375f4c43c0ecdb3fee8713981f2a7b5c535d7d7956846c95556c874d3f33b34ffa9f5ecb8c3c1e4790cfc0455542cdf3c516cc605bc531bd
+  data.tar.gz: 8414902bd41d62597eb6b583d1c5feecfa2c7fc50a04bf055c7ff21249d12ebcb467aad6b8e0aa483fd3435d954bbb1e80e14a36e5523ac656353db9589f31ca

data/.github/workflows/lint.yaml

@@ -0,0 +1,21 @@
+name: 'Lint'
+
+on: push
+
+jobs:
+  run-rubocop:
+    name: Lint code
+    strategy:
+      fail-fast: false
+      matrix:
+        rubyVersion: ["2.7.6"]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: ${{ matrix.rubyVersion }}
+          bundler-cache: true
+          cache-version: "${{ matrix.rubyVersion }}-0"
+      - name: Run RuboCop
+        run: bundle exec rubocop

data/.github/workflows/test.yaml

@@ -0,0 +1,31 @@
+name: Test
+
+on: push
+
+jobs:
+  run-tests:
+    name: Test code
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        ruby:
+          - 3.1.2
+          - 3.0.4
+          - 2.7.6
+    steps:
+      - uses: actions/checkout@v3
+      - uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: ${{ matrix.ruby }}
+          bundler-cache: true
+          cache-version: ${{ matrix.ruby }}-0
+      - name: Run RSpec
+        env:
+          CI: true
+        run: bundle exec rspec
+      - name: Coveralls
+        uses: coverallsapp/github-action@v1.1.2
+        continue-on-error: true
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}

data/.gitignore

@@ -0,0 +1,15 @@
+/.bundle/
+/.yardoc
+/_yardoc/
+/coverage/
+/doc/
+/pkg/
+/spec/reports/
+/tmp/
+
+# rspec failure tracking
+.rspec_status
+.DS_Store
+.vscode
+
+*.gem

data/.ruby-version

@@ -0,0 +1 @@
+3.1.2

data/CHANGELOG.md

@@ -0,0 +1,3 @@
+# 1.0.0
+
+Initial release. Extracted from an internal library.

data/Gemfile

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in gemspec
+gemspec

data/Gemfile.lock

@@ -0,0 +1,114 @@
+PATH
+  remote: .
+  specs:
+    s3_selector (0.1.1)
+      aws-sdk-athena
+      aws-sdk-s3
+      concurrent_executor
+
+GEM
+  remote: https://rubygems.org/
+  specs:
+    addressable (2.8.5)
+      public_suffix (>= 2.0.2, < 6.0)
+    ast (2.4.2)
+    aws-eventstream (1.2.0)
+    aws-partitions (1.797.0)
+    aws-sdk-athena (1.74.0)
+      aws-sdk-core (~> 3, >= 3.177.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-core (3.180.1)
+      aws-eventstream (~> 1, >= 1.0.2)
+      aws-partitions (~> 1, >= 1.651.0)
+      aws-sigv4 (~> 1.5)
+      jmespath (~> 1, >= 1.6.1)
+    aws-sdk-kms (1.71.0)
+      aws-sdk-core (~> 3, >= 3.177.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-s3 (1.132.0)
+      aws-sdk-core (~> 3, >= 3.179.0)
+      aws-sdk-kms (~> 1)
+      aws-sigv4 (~> 1.6)
+    aws-sigv4 (1.6.0)
+      aws-eventstream (~> 1, >= 1.0.2)
+    benchmark-ips (2.12.0)
+    coderay (1.1.3)
+    concurrent_executor (1.3.0)
+    crack (0.4.5)
+      rexml
+    diff-lcs (1.5.0)
+    docile (1.4.0)
+    hashdiff (1.0.1)
+    jmespath (1.6.2)
+    json (2.6.3)
+    language_server-protocol (3.17.0.3)
+    method_source (1.0.0)
+    parallel (1.23.0)
+    parser (3.2.2.3)
+      ast (~> 2.4.1)
+      racc
+    pry (0.14.2)
+      coderay (~> 1.1)
+      method_source (~> 1.0)
+    public_suffix (5.0.3)
+    racc (1.7.1)
+    rainbow (3.1.1)
+    rake (10.5.0)
+    regexp_parser (2.8.1)
+    rexml (3.2.6)
+    rspec (3.12.0)
+      rspec-core (~> 3.12.0)
+      rspec-expectations (~> 3.12.0)
+      rspec-mocks (~> 3.12.0)
+    rspec-core (3.12.2)
+      rspec-support (~> 3.12.0)
+    rspec-expectations (3.12.3)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.12.0)
+    rspec-mocks (3.12.6)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.12.0)
+    rspec-support (3.12.1)
+    rubocop (1.55.1)
+      json (~> 2.3)
+      language_server-protocol (>= 3.17.0)
+      parallel (~> 1.10)
+      parser (>= 3.2.2.3)
+      rainbow (>= 2.2.2, < 4.0)
+      regexp_parser (>= 1.8, < 3.0)
+      rexml (>= 3.2.5, < 4.0)
+      rubocop-ast (>= 1.28.1, < 2.0)
+      ruby-progressbar (~> 1.7)
+      unicode-display_width (>= 2.4.0, < 3.0)
+    rubocop-ast (1.29.0)
+      parser (>= 3.2.1.0)
+    ruby-progressbar (1.13.0)
+    simplecov (0.22.0)
+      docile (~> 1.1)
+      simplecov-html (~> 0.11)
+      simplecov_json_formatter (~> 0.1)
+    simplecov-html (0.12.3)
+    simplecov_json_formatter (0.1.4)
+    unicode-display_width (2.4.2)
+    webmock (3.18.1)
+      addressable (>= 2.8.0)
+      crack (>= 0.3.2)
+      hashdiff (>= 0.4.0, < 2.0.0)
+
+PLATFORMS
+  x86_64-darwin-21
+  x86_64-linux
+
+DEPENDENCIES
+  benchmark-ips
+  bundler (~> 2.0)
+  pry
+  rake (~> 10.0)
+  rspec (~> 3.0)
+  rubocop
+  s3_selector!
+  simplecov
+  webmock
+
+BUNDLED WITH
+   2.4.13

data/README.md

@@ -0,0 +1,23 @@
+# S3 Selector
+
+When selecting large amounts of data via the AWS SDK, it collects all results into memory and then provides you an enumerable. This is bad for very large result sets. This gem solves that by streaming results in, instead.
+
+## Installation
+
+Add this line to your application's Gemfile:
+
+```ruby
+gem 's3_selector'
+```
+
+And then execute:
+
+    $ bundle
+
+Or install it yourself as:
+
+    $ gem install s3_selector
+
+## Usage
+
+### With S3 Files

data/Rakefile

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require 'bundler/gem_tasks'
+require 'rspec/core/rake_task'
+
+RSpec::Core::RakeTask.new(:spec)
+
+task default: :spec

data/bin/console

@@ -0,0 +1,11 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "instrument_all_the_things"
+
+# You can add fixtures and/or initialization code here to make experimenting
+# with your gem easier. You can also use a different console, if you like.
+
+# (If you use this, don't forget to add pry to your Gemfile!)
+require "pry"
+Pry.start

data/bin/setup

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+IFS=$'\n\t'
+set -vx
+
+bundle install
+
+# Do any other automated setup that you need to do here

data/lib/s3_selector/results_stream.rb

@@ -0,0 +1,163 @@
+class ResultsStream
+  DEFAULT_S3_SELECT_THREADS = 5
+  DELIMITER = "\n".bytes.first.freeze
+
+  attr_accessor :s3_files, :number_of_threads, :query, :s3_client, :region,
+                :input_format, :input_format_options
+
+  def initialize(
+    s3_files:,
+    query: 'SELECT * FROM S3Object',
+    s3_client: Aws::S3::Client.new,
+    number_of_threads: DEFAULT_S3_SELECT_THREADS,
+    region: 'us-east-1',
+    input_format: :parquet,
+    input_format_options: {}
+  )
+    self.s3_files             = s3_files
+    self.s3_client            = s3_client
+    self.query                = query
+    self.number_of_threads    = number_of_threads
+    self.region               = region
+    self.input_format         = input_format
+    self.input_format_options = input_format_options
+  end
+
+  def records
+    Enumerator.new do |yielder|
+      if s3_files.length == 1
+        read_s3_file(file: s3_files.first) { |data| yielder << data }
+      else
+        ConcurrentExecutor.consume_enumerable(s3_files) do |s3_file|
+          read_s3_file(file: s3_file) { |data| yielder << data }
+        end
+      end
+    end
+  end
+
+  private
+
+  def signer
+    @signer ||= Aws::Sigv4::Signer.new(
+      service: 's3',
+      region: region,
+      credentials_provider: Aws::CredentialProviderChain.new.resolve,
+      uri_escape_path: false,
+    )
+  end
+
+  def s3_select_request_body
+    @s3_select_request_body ||= <<~XML
+      <SelectObjectContentRequest xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
+        <Expression>#{query}</Expression>
+        <ExpressionType>SQL</ExpressionType>
+        <InputSerialization>
+          #{input_serialization_body}
+        </InputSerialization>
+        <OutputSerialization>
+          <JSON>
+            <RecordDelimiter>\n</RecordDelimiter>
+          </JSON>
+        </OutputSerialization>
+      </SelectObjectContentRequest>
+    XML
+  end
+
+  def input_serialization_body
+    case input_format
+    when :parquet then '<Parquet />'
+    when :json then json_input_serialization_body
+    else
+      raise "unknown input format #{input_format}"
+    end
+  end
+
+  def json_input_serialization_body
+    format = input_format_options&.fetch(:type, '')&.to_s&.upcase
+
+    raise 'When querying JSON you must specify a input format option `:type` of either DOCUMENT or LINES' unless %w[DOCUMENT LINES].include?(format)
+
+    compression = input_format_options&.fetch(:compression, '')&.to_s&.upcase
+
+    <<~BODY
+      #{compression && compression != '' ? "<CompressionType>#{compression}</CompressionType>" : ''}
+      <JSON>
+         <Type>#{format}</Type>
+      </JSON>
+    BODY
+  end
+
+  def decoder
+    # Do NOT memoize this. Need a fresh stream decoder per HTTP request
+    Aws::EventStream::Decoder.new
+  end
+
+  def bytes_to_utf8(bytes)
+    bytes.pack('C*').force_encoding('utf-8')
+  end
+
+  def split_and_ensure_delimiter(buffer)
+    found = []
+    position = 0
+    last_slice_end = 0
+
+    while position < buffer.length
+      if buffer[position] == DELIMITER
+        found << buffer[last_slice_end..(position - 1)]
+        last_slice_end = position += 1
+      else
+        position += 1
+      end
+    end
+
+    [found, buffer[last_slice_end..]]
+  end
+
+  def read_s3_file(file:)
+    file_request_decoder = decoder
+
+    url = URI("#{file.public_url}?select&select-type=2")
+
+    signature = signer.sign_request(
+      http_method: 'POST',
+      url: url,
+      headers: { 'Content-Type' => 'application/octet-stream' },
+      body: s3_select_request_body,
+    )
+
+    request = Net::HTTP::Post.new(url.request_uri, signature.headers.merge({ 'Content-Type' => 'application/octet-stream' }))
+    request.body = s3_select_request_body
+
+    mon = Monitor.new
+
+    Net::HTTP.start(url.host, url.port, use_ssl: true) do |http|
+      bytes = []
+      http.request(request) do |response|
+        response.read_body do |chunk|
+          message, chunk_eof = file_request_decoder.decode_chunk(chunk)
+          if message && message.headers[':event-type'].value == 'Records'
+            lines = []
+
+            mon.synchronize do
+              bytes.concat(message.payload.read.bytes)
+              lines, bytes = split_and_ensure_delimiter(bytes)
+            end
+
+            lines.each do |line|
+              yield JSON.parse(bytes_to_utf8(line))
+            rescue JSON::ParserError
+            end
+          end
+
+          if chunk_eof && !!bytes && !bytes.empty?
+            begin
+              yield JSON.parse(bytes_to_utf8(bytes))
+              bytes = []
+            rescue JSON::ParserError
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/s3_selector/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module S3Selector
+  VERSION = '0.1.1'
+end

data/lib/s3_selector.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require 's3_selector/version'
+require 'aws-sdk-s3'
+require 'concurrent_executor'
+
+module S3Selector
+end

data/s3_selector.gemspec

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+lib = File.expand_path('lib', __dir__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 's3_selector/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = 's3_selector'
+  spec.version       = S3Selector::VERSION
+  spec.authors       = ['Nathaniel Rowe']
+  spec.email         = ['nathaniel.rowe@terminus.com']
+
+  spec.summary       = 'Stream s3 select results'
+  spec.description   = 'Stream s3 select results'
+  spec.homepage      = 'https://github.com/GetTerminus/s3_selector'
+
+  spec.metadata['allowed_push_host'] = 'https://www.rubygems.org'
+
+  spec.metadata['homepage_uri'] = spec.homepage
+  spec.metadata['source_code_uri'] = 'https://github.com/GetTerminus/s3_selector'
+  # spec.metadata['changelog_uri'] = 'http://google.com'
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files = Dir.chdir(File.expand_path(__dir__)) do
+    `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features|Gemfile.lock|vendor)/}) }
+  end
+  spec.bindir        = 'exe'
+  spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
+  spec.require_paths = ['lib']
+  spec.required_ruby_version = Gem::Requirement.new('>= 2.7.0')
+
+  spec.add_dependency 'aws-sdk-s3'
+  spec.add_dependency 'aws-sdk-athena'
+  spec.add_dependency 'concurrent_executor'
+
+
+  spec.add_development_dependency 'benchmark-ips'
+  spec.add_development_dependency 'bundler', '~> 2.0'
+  spec.add_development_dependency 'pry'
+  spec.add_development_dependency 'rake', '~> 10.0'
+  spec.add_development_dependency 'rspec', '~> 3.0'
+  spec.add_development_dependency 'rubocop'
+  spec.add_development_dependency 'simplecov'
+  spec.add_development_dependency 'webmock'
+end

metadata

@@ -0,0 +1,214 @@
+--- !ruby/object:Gem::Specification
+name: s3_selector
+version: !ruby/object:Gem::Version
+  version: 0.1.1
+platform: ruby
+authors:
+- Nathaniel Rowe
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2023-08-04 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-s3
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-athena
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: concurrent_executor
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: benchmark-ips
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: bundler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.0'
+- !ruby/object:Gem::Dependency
+  name: pry
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: rubocop
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: simplecov
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: webmock
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description: Stream s3 select results
+email:
+- nathaniel.rowe@terminus.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".github/workflows/lint.yaml"
+- ".github/workflows/test.yaml"
+- ".gitignore"
+- ".ruby-version"
+- CHANGELOG.md
+- Gemfile
+- Gemfile.lock
+- README.md
+- Rakefile
+- bin/console
+- bin/setup
+- lib/s3_selector.rb
+- lib/s3_selector/results_stream.rb
+- lib/s3_selector/version.rb
+- s3_selector.gemspec
+homepage: https://github.com/GetTerminus/s3_selector
+licenses: []
+metadata:
+  allowed_push_host: https://www.rubygems.org
+  homepage_uri: https://github.com/GetTerminus/s3_selector
+  source_code_uri: https://github.com/GetTerminus/s3_selector
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 2.7.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.3.12
+signing_key:
+specification_version: 4
+summary: Stream s3 select results
+test_files: []