feat: SelectObjectContentRequest

lib/ex_aws/s3.ex

@@ -757,6 +757,113 @@ defmodule ExAws.S3 do
     }
   end
 
+  @type select_object_content_opts :: [
+          {:input_serialization, input_serialization}
+          | {:output_serialization, output_serialization}
+          | {:scan_range, scan_range}
+        ]
+
+  @type input_serialization ::
+          %{csv_input: csv_input} | %{json_input: json_input} | %{parquet_input: %{}}
+
+  @type csv_input :: %{
+          file_header_info: :none | :ignore | :use,
+          comments: binary,
+          quote_escape_character: binary,
+          record_delimiter: binary,
+          field_delimiter: binary,
+          quote_character: binary,
+          allow_quoted_record_delimiter: boolean
+        }
+  @type json_input :: %{
+          type: :document | :lines
+        }
+
+  @type output_serialization ::
+          %{csv_output: csv_output} | %{json_output: json_output}
+
+  @type csv_output :: %{
+          quote_fields: :always | :as_needed,
+          quote_escape_character: binary,
+          record_delimiter: binary,
+          field_delimiter: binary,
+          quote_character: binary
+        }
+
+  @type json_output :: %{
+          record_delimiter: binary
+        }
+
+  @type scan_range :: %{start: pos_integer, end: pos_integer}
+
+  @doc """
+  Filters and selects the contents of an Amazon S3 object based on an SQL statement.
+
+  ## Options
+  * `:input_serialization` - Specifies JSON, CSV, or Parquet as the input serialization format.
+     each of which has corresponding parameters to describe the format of the object to be retrieved.
+  * `:output_serialization` - Specifies JSON or CSV as the output serialization format.
+     Each of which has corresponding parameters to describe the format of the output data.
+  * `:scan_range` - Specifies the byte range of the object to get the records from.
+
+  More information can be found in the [AWS documentation](https://docs.aws.amazon.com/AmazonS3/latest/API/API_SelectObjectContent.html).
+
+
+  ## Example
+  ```
+  S3.select_object_content(
+    "my-bucket",
+    "path/to/file.csv",
+    "SELECT * FROM S3Object s WHERE s._1 = 'some value'",
+    input_serialization: %{
+      csv_input: %{
+        file_header_info: :use,
+        record_delimiter: "\\n",
+        field_delimiter: ",",
+        quote_character: "\\""
+      }
+    },
+    output_serialization: %{
+      csv_output: %{
+        record_delimiter: "\\n",
+        field_delimiter: ",",
+        quote_character: "\\""
+      }
+    }
+  ) |> ExAws.stream!()
+  ```
+
+  Note that **this won't start fetching anything immediately** since it returns an Elixir `Stream`.
+
+  ### Streaming into a file
+  ```
+  S3.select_object_content(
+    "my-bucket",
+    "path/to/file.csv",
+    "SELECT * FROM S3Object s WHERE s._1 = 'some value'"
+  ) |> ExAws.stream!() |> Stream.into(File.stream!("output.csv"))
+  ```
+  """
+  @spec select_object_content(
+          bucket :: binary,
+          path :: binary,
+          query :: binary,
+          opts :: select_object_content_opts
+        ) :: __MODULE__.SelectObjectContent.t()
+  def select_object_content(
+        bucket,
+        path,
+        query,
+        opts \\ []
+      ) do
+    %__MODULE__.SelectObjectContent{
+      bucket: bucket,
+      path: path,
+      query: query,
+      opts: opts
+    }
+  end
+
   @type upload_opt ::
           {:max_concurrency, pos_integer}
           | {:timeout, pos_integer}

lib/ex_aws/s3/parsers/event_stream.ex

@@ -0,0 +1,38 @@
+defmodule ExAws.S3.Parsers.EventStream do
+  @moduledoc false
+
+  # Parses EventStream messages.
+
+  # AWS encodes EventStream messages in binary as follows:
+  # [      prelude     ][     headers   ][    payload    ][   message-crc  ]
+  # |<--  12 bytes  -->|<-- variable -->|<-- variable -->|<--  4 bytes  -->|
+
+  # This module parses this information and returns a struct with the prelude, headers and payload.
+  # The prelude contains the total length of the message, the length of the headers,
+  # the length of the prelude, the CRC of the message, and the length of the payload.
+
+  # The headers are a map of header names to values.
+  # The payload is the actual message data.
+  # The message-crc is a CRC32 checksum of the message (excluding the message-crc itself).
+  # Refer to https://docs.aws.amazon.com/AmazonS3/latest/API/RESTSelectObjectAppendix.html for more information.
+
+  alias ExAws.S3.Parsers.EventStream.Message
+
+  defp parse_message(chunk) do
+    with {:ok, message} <- Message.parse(chunk) do
+      message
+    end
+  end
+
+  def parse_raw_stream(
+        {:ok,
+         %{
+           stream: stream
+         }}
+      ) do
+    stream
+    |> Stream.map(&parse_message/1)
+    |> Stream.filter(&Message.is_record?/1)
+    |> Stream.map(&Message.get_payload/1)
+  end
+end

lib/ex_aws/s3/parsers/event_stream/header.ex

@@ -0,0 +1,52 @@
+defmodule ExAws.S3.Parsers.EventStream.Header do
+  @moduledoc false
+
+  # Parses EventStream headers.
+
+  # AWS encodes EventStream headers as follows:
+
+  # [header-name-size][header-name][header-data-type][header-value-size][header-value-data]
+  # |<--  1 byte  -->|<-variable->|<--   1 byte  -->|<--  2 bytes   -->|<--  variable  -->|
+
+  # This module parses this information and returns a map of header names - values.
+  # header-data-type is always 0x07(String) for S3.
+
+  alias ExAws.S3.Parsers.EventStream.Prelude
+
+  def extract_headers(header_bytes) do
+    do_extract_headers(header_bytes, [])
+  end
+
+  defp do_extract_headers(<<>>, headers), do: Map.new(headers)
+
+  defp do_extract_headers(<<header_name_size, rest::binary>>, headers) do
+    <<header_name::binary-size(header_name_size), <<7>>, rest::binary>> = rest
+    <<value_size_binary::binary-size(2), rest::binary>> = rest
+
+    value_size = :binary.decode_unsigned(value_size_binary, :big)
+    <<value::binary-size(value_size), rest::binary>> = rest
+
+    do_extract_headers(rest, [{header_name, value} | headers])
+  end
+
+  defp extract_header_bytes(
+         %Prelude{
+           prelude_length: prelude_length,
+           headers_length: headers_length
+         },
+         payload_bytes
+       ) do
+    <<_prelude::binary-size(prelude_length), headers_bytes::binary-size(headers_length),
+      _payload::binary>> = payload_bytes
+
+    headers_bytes
+  end
+
+  def parse(
+        %Prelude{} = prelude,
+        payload_bytes
+      ) do
+    headers = prelude |> extract_header_bytes(payload_bytes) |> extract_headers()
+    {:ok, headers}
+  end
+end

lib/ex_aws/s3/parsers/event_stream/message.ex

@@ -0,0 +1,67 @@
+defmodule ExAws.S3.Parsers.EventStream.Message do
+  @moduledoc false
+
+  # Parses EventStream messages. This module parses this information and returns
+  # a struct with the prelude, headers and payload. Also verifies the message CRC.
+
+  alias ExAws.S3.Parsers.EventStream.Prelude
+  alias ExAws.S3.Parsers.EventStream.Header
+  import Bitwise
+
+  defstruct prelude: nil,
+            headers: nil,
+            payload: nil
+
+  def verify_message_crc(
+        %Prelude{
+          crc: prelude_crc,
+          payload_length: payload_length,
+          headers_length: headers_length
+        },
+        payload_bytes
+      ) do
+    <<_::binary-size(8), message_bytes::binary-size(payload_length + headers_length + 4),
+      message_crc_bytes::binary-size(4)>> = payload_bytes
+
+    message_crc = :binary.decode_unsigned(message_crc_bytes, :big)
+    computed_crc = prelude_crc |> :erlang.crc32(message_bytes) |> band(0xFFFFFFFF)
+
+    if computed_crc == message_crc do
+      :ok
+    else
+      {:error, :message_checksum_mismatch}
+    end
+  end
+
+  def parse_payload(
+        %Prelude{
+          prelude_length: prelude_length,
+          headers_length: headers_length,
+          payload_length: payload_length
+        },
+        payload_bytes
+      ) do
+    <<_prelude_headers::binary-size(prelude_length + headers_length),
+      message_bytes::binary-size(payload_length), _::binary>> = payload_bytes
+
+    {:ok, message_bytes}
+  end
+
+  def is_record?(%__MODULE__{headers: headers}) do
+    Map.get(headers, ":event-type") == "Records"
+  end
+
+  def get_payload(%__MODULE__{payload: payload}) do
+    payload
+  end
+
+  def parse(chunk) do
+    with {:ok, prelude, payload_bytes} <-
+           Prelude.parse(chunk),
+         :ok <- verify_message_crc(prelude, payload_bytes),
+         {:ok, headers} <- Header.parse(prelude, payload_bytes),
+         {:ok, payload} <- parse_payload(prelude, payload_bytes) do
+      {:ok, %__MODULE__{prelude: prelude, payload: payload, headers: headers}}
+    end
+  end
+end

lib/ex_aws/s3/parsers/event_stream/prelude.ex

@@ -0,0 +1,88 @@
+defmodule ExAws.S3.Parsers.EventStream.Prelude do
+  @moduledoc false
+  # Parses EventStream preludes.
+
+  # AWS encodes EventStream preludes in binary as follows:
+
+  # [  total-length  ][headers-length][  prelude crc ]
+  # |<--  4 bytes -->|<-- 4 bytes -->|<-- 4 bytes -->|
+
+  # This module parses this information and returns a struct with the total length of the message,
+  # the length of the headers, the length of the prelude, the CRC of the message,
+  # and the length of the payload.
+  import Bitwise
+
+  defstruct total_length: nil,
+            headers_length: nil,
+            prelude_length: nil,
+            crc: nil,
+            payload_length: nil,
+            prelude_bytes: nil
+
+  @prelude_length 12
+  # 128 Kb
+  @max_header_length 128 * 1024
+  # 16 Mb
+  @max_payload_length 16 * 1024 * 1024
+
+  defp unpack_prelude(
+         <<
+           total_length_bytes::binary-size(4),
+           headers_length_bytes::binary-size(4),
+           crc::binary-size(4)
+         >> = prelude_bytes
+       ) do
+    total_length = :binary.decode_unsigned(total_length_bytes, :big)
+    headers_length = :binary.decode_unsigned(headers_length_bytes, :big)
+    crc = :binary.decode_unsigned(crc, :big)
+
+    {:ok,
+     %__MODULE__{
+       total_length: total_length,
+       headers_length: headers_length,
+       prelude_length: @prelude_length,
+       payload_length: total_length - @prelude_length - headers_length - 4,
+       crc: crc,
+       prelude_bytes: prelude_bytes
+     }}
+  end
+
+  def validate_prelude(
+        %__MODULE__{
+          headers_length: headers_length,
+          payload_length: payload_length
+        } = prelude
+      ) do
+    cond do
+      headers_length > @max_header_length ->
+        {:error, :invalid_headers_length}
+
+      payload_length > @max_payload_length ->
+        {:error, :invalid_payload_length}
+
+      true ->
+        {:ok, prelude}
+    end
+  end
+
+  def validate_checksum(
+        <<prelude_bytes_without_crc::binary-size(@prelude_length - 4), _rest::binary>>,
+        prelude_checksum
+      ) do
+    computed_checksum = prelude_bytes_without_crc |> :erlang.crc32() |> band(0xFFFFFFFF)
+
+    if computed_checksum == prelude_checksum do
+      :ok
+    else
+      {:error, :prelude_checksum_mismatch}
+    end
+  end
+
+  def parse(<<prelude_bytes::binary-size(@prelude_length), _rest::binary>> = payload) do
+    with {:ok, unpacked_prelude} <- unpack_prelude(prelude_bytes),
+         {:ok, prelude} <- validate_prelude(unpacked_prelude),
+         :ok <- validate_checksum(prelude_bytes, unpacked_prelude.crc) do
+      {:ok, prelude, payload}
+    end
+  end
+end