| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| class FlightSQLTest < Test::Unit::TestCase |
| include Helper::Sandbox |
| |
| def max_n_rows_per_record_batch |
| 10 |
| end |
| |
| setup do |
| unless flight_client.respond_to?(:authenticate_basic) |
| omit("ArrowFlight::Client#authenticate_basic is needed" ) |
| end |
| @options = ArrowFlight::CallOptions.new |
| @options.add_header("x-flight-sql-database", @test_db_name) |
| user = @postgresql.user |
| password = @postgresql.password |
| flight_client.authenticate_basic(user, password, @options) |
| end |
| |
| def to_sql(value) |
| case value |
| when Time |
| "'#{value.dup.utc.strftime("%Y-%m-%d %H:%M:%S.%6N")}'" |
| when String |
| sql_string = "'" |
| value.each_char do |char| |
| case char |
| when "'" |
| sql_string << "''" |
| when "\\" |
| sql_string << "\\\\" |
| else |
| if (0..31).cover?(char.ord) or (127..255).cover?(char.ord) |
| sql_string << ("\\%03d" % char.ord) |
| else |
| sql_string << char |
| end |
| end |
| end |
| sql_string << "'" |
| sql_string |
| else |
| value.to_s |
| end |
| end |
| |
| timestamp_value = Time.parse("2023-08-02T18:03:13.526572321Z").utc |
| |
| data("int16", ["smallint", :int16, -2]) |
| data("int32", ["integer", :int32, -2]) |
| data("int64", ["bigint", :int64, -2]) |
| data("float", ["real", :float, -2.2]) |
| data("double", ["double precision", :double, -2.2]) |
| data("string - text", ["text", :string, "b"]) |
| data("string - varchar", ["varchar(10)", :string, "b"]) |
| data("binary", ["bytea", :binary, "\x0".b]) |
| data("timestamp", ["timestamp", [:timestamp, :micro], timestamp_value]) |
| def test_select_direct |
| pg_type, data_type, value = data |
| data_type = Arrow::DataType.resolve(data_type) |
| values = data_type.build_array([value]) |
| sql = "SELECT #{to_sql(value)}::#{pg_type} AS value" |
| info = flight_sql_client.execute(sql, @options) |
| assert_equal(Arrow::Schema.new(value: values.value_data_type), |
| info.get_schema) |
| endpoint = info.endpoints.first |
| reader = flight_sql_client.do_get(endpoint.ticket, @options) |
| assert_equal(Arrow::Table.new(value: values), |
| reader.read_all) |
| end |
| |
| data("int16", [:int16, -2]) |
| data("int32", [:int32, -2]) |
| data("int64", [:int64, -2]) |
| data("float", [:float, -2.2]) |
| data("double", [:double, -2.2]) |
| data("string", [:string, "b"]) |
| data("binary", [:binary, "\x0".b]) |
| data("timestamp", [[:timestamp, :micro], timestamp_value]) |
| def test_select_prepare |
| unless flight_sql_client.respond_to?(:prepare) |
| omit("red-arrow-flight-sql 14.0.0 or later is required") |
| end |
| |
| data_type, value = data |
| data_type = Arrow::DataType.resolve(data_type) |
| values = data_type.build_array([value]) |
| flight_sql_client.prepare("SELECT $1 AS value", |
| @options) do |statement| |
| statement.record_batch = Arrow::RecordBatch.new(value: values) |
| info = statement.execute(@options) |
| assert_equal(Arrow::Schema.new(value: values.value_data_type), |
| info.get_schema) |
| endpoint = info.endpoints.first |
| reader = flight_sql_client.do_get(endpoint.ticket, @options) |
| assert_equal(Arrow::Table.new(value: values), |
| reader.read_all) |
| end |
| end |
| |
| def test_select_prepare_without_parameters |
| unless flight_sql_client.respond_to?(:prepare) |
| omit("red-arrow-flight-sql 14.0.0 or later is required") |
| end |
| |
| flight_sql_client.prepare("SELECT 29 AS value", @options) do |statement| |
| info = statement.execute(@options) |
| assert_equal(Arrow::Schema.new(value: :int32), |
| info.get_schema) |
| endpoint = info.endpoints.first |
| reader = flight_sql_client.do_get(endpoint.ticket, @options) |
| assert_equal(Arrow::Table.new(value: Arrow::Int32Array.new([29])), |
| reader.read_all) |
| end |
| end |
| |
| def test_select_from |
| run_sql("CREATE TABLE data (value integer)") |
| run_sql("INSERT INTO data VALUES (1), (-2), (3)") |
| |
| info = flight_sql_client.execute("SELECT * FROM data", @options) |
| assert_equal(Arrow::Schema.new(value: :int32), |
| info.get_schema) |
| endpoint = info.endpoints.first |
| reader = flight_sql_client.do_get(endpoint.ticket, @options) |
| assert_equal(Arrow::Table.new(value: Arrow::Int32Array.new([1, -2, 3])), |
| reader.read_all) |
| end |
| |
| data(:type, ["integer", "text"]) |
| data(:may_null, [true, false]) |
| def test_select_multiple_record_batches_response |
| type = data[:type] |
| may_null = data[:may_null] |
| if type == "integer" |
| arrow_type = Arrow::Int32DataType.new |
| values = 30.times.to_a |
| else |
| arrow_type = Arrow::StringDataType.new |
| values = 30.times.collect(&:to_s) |
| end |
| |
| create_table = "CREATE TABLE data (value #{type}" |
| # XXX: "NOT NULL" doesn't change FormData_pg_attribute::attnotnull |
| # returned by "SELECT *". So the "may_null" data driven test |
| # parameter is meaningless... |
| create_table << " NOT NULL" unless may_null |
| create_table << ")" |
| run_sql(create_table) |
| run_sql("INSERT INTO data VALUES " + |
| values.collect {|value| "(#{value})"}.join(", ")) |
| |
| info = flight_sql_client.execute("SELECT * FROM data", @options) |
| assert_equal(Arrow::Schema.new(value: arrow_type), |
| info.get_schema) |
| endpoint = info.endpoints.first |
| reader = flight_sql_client.do_get(endpoint.ticket, @options) |
| chunks = values.each_slice(max_n_rows_per_record_batch).collect do |chunk| |
| arrow_type.array_class.new(chunk) |
| end |
| assert_equal(Arrow::Table.new(value: Arrow::ChunkedArray.new(chunks)), |
| reader.read_all) |
| end |
| |
| def test_insert_direct |
| unless flight_sql_client.respond_to?(:execute_update) |
| omit("red-arrow-flight-sql 13.0.0 or later is required") |
| end |
| |
| run_sql("CREATE TABLE data (value integer)") |
| |
| n_changed_records = flight_sql_client.execute_update( |
| "INSERT INTO data VALUES (1), (-2), (3)", |
| @options) |
| assert_equal(3, n_changed_records) |
| assert_equal([<<-RESULT, ""], run_sql("SELECT * FROM data")) |
| SELECT * FROM data |
| value |
| ------- |
| 1 |
| -2 |
| 3 |
| (3 rows) |
| |
| RESULT |
| end |
| |
| timestamp_values = [ |
| Time.parse("2023-08-02T18:03:13.526572321Z").utc, |
| Time.parse("2023-08-02T18:04:13.526572321Z").utc, |
| Time.parse("2023-08-02T18:05:13.526572321Z").utc, |
| ] |
| |
| data("int8", ["smallint", :int8, [1, -2, 3]]) |
| data("int16", ["smallint", :int16, [1, -2, 3]]) |
| data("int32", ["integer", :int32, [1, -2, 3]]) |
| data("int64", ["bigint", :int64, [1, -2, 3]]) |
| data("uint8", ["smallint", :uint8, [1, 2, 3]]) |
| data("uint16", ["smallint", :uint16, [1, 2, 3]]) |
| data("uint32", ["integer", :uint32, [1, 2, 3]]) |
| data("uint64", ["bigint", :uint64, [1, 2, 3]]) |
| data("float", ["real", :float, [1.1, -2.2, 3.3]]) |
| data("double", ["double precision", :double, [1.1, -2.2, 3.3]]) |
| data("string - text", ["text", :string, ["a", "b", "c"]]) |
| data("string - varchar", ["varchar(10)", :string, ["a", "b", "c"]]) |
| data("binary", ["bytea", :binary, ["\x0".b, "\x1".b, "\x2".b]]) |
| data("timestamp(second)", |
| ["timestamp", [:timestamp, :second], timestamp_values]) |
| data("timestamp(milli)", |
| ["timestamp", [:timestamp, :milli], timestamp_values]) |
| data("timestamp(micro)", |
| ["timestamp", [:timestamp, :micro], timestamp_values]) |
| data("timestamp(nano)", |
| ["timestamp", [:timestamp, :nano], timestamp_values]) |
| def test_insert_prepare |
| unless flight_sql_client.respond_to?(:prepare) |
| omit("red-arrow-flight-sql 14.0.0 or later is required") |
| end |
| |
| pg_type, data_type, values = data |
| data_type = Arrow::DataType.resolve(data_type) |
| run_sql("CREATE TABLE data (value #{pg_type})") |
| |
| array = data_type.build_array(values) |
| flight_sql_client.prepare("INSERT INTO data VALUES ($1)", |
| @options) do |statement| |
| statement.record_batch = Arrow::RecordBatch.new(value: array) |
| n_changed_records = statement.execute_update(@options) |
| assert_equal(values.size, n_changed_records) |
| end |
| |
| case values.first |
| when Time |
| case data_type.unit |
| when Arrow::TimeUnit::SECOND |
| value_size = "1970-01-01 00:00:00".size |
| strftime_format = "%Y-%m-%d %H:%M:%S" |
| when Arrow::TimeUnit::MILLI |
| value_size = "1970-01-01 00:00:00.000".size |
| strftime_format = "%Y-%m-%d %H:%M:%S.%3N" |
| when Arrow::TimeUnit::MICRO, Arrow::TimeUnit::NANO |
| value_size = "1970-01-01 00:00:00.000000".size |
| strftime_format = "%Y-%m-%d %H:%M:%S.%6N" |
| else |
| raise "unsupported: #{data_type.unit.inspect}" |
| end |
| else |
| value_size = 5 |
| end |
| output = <<-RESULT |
| SELECT * FROM data |
| #{"value".center(value_size)} |
| -#{"-" * value_size}- |
| RESULT |
| values.each do |value| |
| case value |
| when Float |
| output << (" %5.1f\n" % value) |
| when Integer |
| output << (" %5d\n" % value) |
| when Time |
| output << " #{value.strftime(strftime_format)}\n" |
| else |
| if value.encoding == "".b.encoding |
| output << " " |
| value.each_byte do |byte| |
| output << ("\\x%02x" % byte) |
| end |
| output << "\n" |
| else |
| output << " #{value}\n" |
| end |
| end |
| end |
| output << "(#{values.size} rows)\n" |
| output << "\n" |
| assert_equal([output, ""], run_sql("SELECT * FROM data")) |
| end |
| end |
