From 3d7f99f9500d291474c6b4ad4c99a02c6f854fd6 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 5 Aug 2022 12:55:00 -0400 Subject: [PATCH 01/23] ARROW-17323: [Go] Cleanup and upgrade dependencies --- go/go.mod | 28 ++++++++++++++-------------- go/go.sum | 23 +++++++++++++++++++++-- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/go/go.mod b/go/go.mod index 18a1c9abfa773..0b5a5cf1a88fd 100644 --- a/go/go.mod +++ b/go/go.mod @@ -28,20 +28,20 @@ require ( github.com/golang/snappy v0.0.4 github.com/google/flatbuffers v2.0.6+incompatible github.com/klauspost/asmfmt v1.3.2 - github.com/klauspost/compress v1.15.9 - github.com/kr/pretty v0.1.0 // indirect + github.com/klauspost/compress v1.14.2 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 - github.com/pierrec/lz4/v4 v4.1.15 - github.com/stretchr/testify v1.8.0 - github.com/zeebo/xxh3 v1.0.2 - golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e - golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4 - golang.org/x/sys v0.0.0-20220804214406-8e32c043e418 - golang.org/x/tools v0.1.12 - golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f - gonum.org/v1/gonum v0.11.0 - google.golang.org/grpc v1.48.0 - google.golang.org/protobuf v1.28.1 - gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect + github.com/pierrec/lz4/v4 v4.1.12 + github.com/stretchr/testify v1.7.2 + github.com/zeebo/xxh3 v1.0.1 + golang.org/x/exp v0.0.0-20211216164055-b2b84827b756 + golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect + golang.org/x/sync v0.0.0-20220513210516-0976fa681c29 + golang.org/x/sys v0.0.0-20220412211240-33da011f77ad + golang.org/x/tools v0.1.10 + golang.org/x/xerrors v0.0.0-20220411194840-2f41105eb62f + gonum.org/v1/gonum v0.9.3 + google.golang.org/genproto v0.0.0-20220126215142-9970aeb2e350 // indirect + google.golang.org/grpc v1.44.0 + google.golang.org/protobuf v1.27.1 ) diff --git a/go/go.sum b/go/go.sum index 7653c2cc1b917..1f3bd0392f3f8 100644 --- a/go/go.sum +++ b/go/go.sum @@ -96,11 +96,22 @@ github.com/klauspost/compress v1.15.9 h1:wKRjX6JRtDdrE9qwa4b/Cip7ACOshUI4smpCQan github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU= github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc= github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/lightstep/lightstep-tracer-common/golang/gogo v0.0.0-20190605223551-bc2310a04743/go.mod h1:qklhhLq1aX+mtWk9cPHPzaBjWImj5ULL6C7HFJtXQMM= +github.com/lightstep/lightstep-tracer-go v0.18.1/go.mod h1:jlF1pusYV4pidLvZ+XD0UBX0ZE6WURAspgAczcDHrL4= +github.com/lyft/protoc-gen-validate v0.0.13/go.mod h1:XbGvPuh87YZc5TdIa2/I4pLk0QoUACkjt2znoq26NVQ= +github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= +github.com/mattn/go-isatty v0.0.3/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= +github.com/mattn/go-isatty v0.0.4/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= +github.com/mattn/go-runewidth v0.0.2/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU= +github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= +github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI= @@ -274,11 +285,19 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= -google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= -google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/cheggaaa/pb.v1 v1.0.25/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw= +gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/gcfg.v1 v1.2.3/go.mod h1:yesOnuUOFQAhST5vPY4nbZsb/huCgGGXlipJsBn0b3o= +gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/warnings.v0 v0.1.2/go.mod h1:jksf8JmL6Qr/oQM2OXTHunEvvTAsrWBLb6OOjuVWRNI= +gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bli9HhUf9+ttbYbLASfIpnQbh74= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= From 4d36881b3d86153b8e446a24e451f5b589ce14e9 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 5 Aug 2022 13:15:31 -0400 Subject: [PATCH 02/23] disallow mmap on windows so we can use updated exp pkg. --- go/go.mod | 28 ++++++++++++++-------------- go/go.sum | 23 ++--------------------- 2 files changed, 16 insertions(+), 35 deletions(-) diff --git a/go/go.mod b/go/go.mod index 0b5a5cf1a88fd..18a1c9abfa773 100644 --- a/go/go.mod +++ b/go/go.mod @@ -28,20 +28,20 @@ require ( github.com/golang/snappy v0.0.4 github.com/google/flatbuffers v2.0.6+incompatible github.com/klauspost/asmfmt v1.3.2 - github.com/klauspost/compress v1.14.2 + github.com/klauspost/compress v1.15.9 + github.com/kr/pretty v0.1.0 // indirect github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 - github.com/pierrec/lz4/v4 v4.1.12 - github.com/stretchr/testify v1.7.2 - github.com/zeebo/xxh3 v1.0.1 - golang.org/x/exp v0.0.0-20211216164055-b2b84827b756 - golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect - golang.org/x/sync v0.0.0-20220513210516-0976fa681c29 - golang.org/x/sys v0.0.0-20220412211240-33da011f77ad - golang.org/x/tools v0.1.10 - golang.org/x/xerrors v0.0.0-20220411194840-2f41105eb62f - gonum.org/v1/gonum v0.9.3 - google.golang.org/genproto v0.0.0-20220126215142-9970aeb2e350 // indirect - google.golang.org/grpc v1.44.0 - google.golang.org/protobuf v1.27.1 + github.com/pierrec/lz4/v4 v4.1.15 + github.com/stretchr/testify v1.8.0 + github.com/zeebo/xxh3 v1.0.2 + golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e + golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4 + golang.org/x/sys v0.0.0-20220804214406-8e32c043e418 + golang.org/x/tools v0.1.12 + golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f + gonum.org/v1/gonum v0.11.0 + google.golang.org/grpc v1.48.0 + google.golang.org/protobuf v1.28.1 + gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect ) diff --git a/go/go.sum b/go/go.sum index 1f3bd0392f3f8..7653c2cc1b917 100644 --- a/go/go.sum +++ b/go/go.sum @@ -96,22 +96,11 @@ github.com/klauspost/compress v1.15.9 h1:wKRjX6JRtDdrE9qwa4b/Cip7ACOshUI4smpCQan github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU= github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= -github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= -github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc= github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= -github.com/lightstep/lightstep-tracer-common/golang/gogo v0.0.0-20190605223551-bc2310a04743/go.mod h1:qklhhLq1aX+mtWk9cPHPzaBjWImj5ULL6C7HFJtXQMM= -github.com/lightstep/lightstep-tracer-go v0.18.1/go.mod h1:jlF1pusYV4pidLvZ+XD0UBX0ZE6WURAspgAczcDHrL4= -github.com/lyft/protoc-gen-validate v0.0.13/go.mod h1:XbGvPuh87YZc5TdIa2/I4pLk0QoUACkjt2znoq26NVQ= -github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= -github.com/mattn/go-isatty v0.0.3/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= -github.com/mattn/go-isatty v0.0.4/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= -github.com/mattn/go-runewidth v0.0.2/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU= -github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= -github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI= @@ -285,19 +274,11 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= -gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw= +google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= +google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= -gopkg.in/cheggaaa/pb.v1 v1.0.25/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw= -gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= -gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= -gopkg.in/gcfg.v1 v1.2.3/go.mod h1:yesOnuUOFQAhST5vPY4nbZsb/huCgGGXlipJsBn0b3o= -gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo= -gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= -gopkg.in/warnings.v0 v0.1.2/go.mod h1:jksf8JmL6Qr/oQM2OXTHunEvvTAsrWBLb6OOjuVWRNI= -gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bli9HhUf9+ttbYbLASfIpnQbh74= -gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= From 417037bdcdf35f81b00756b483428fbd99cfece2 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 5 Aug 2022 16:45:09 -0400 Subject: [PATCH 03/23] initial flightsql client --- go/arrow/flight/flightsql/client.go | 96 ++++++++++++++++ go/arrow/flight/flightsql/types.go | 171 +++------------------------- 2 files changed, 110 insertions(+), 157 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index 2f57d05484f76..b4a9bbf970e72 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -19,7 +19,10 @@ package flightsql import ( "context" "errors" +<<<<<<< HEAD "io" +======= +>>>>>>> 7c0a7b52f (initial flightsql client) "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/array" @@ -32,6 +35,7 @@ import ( "google.golang.org/protobuf/types/known/anypb" ) +<<<<<<< HEAD // NewClient is a convenience function to automatically construct // a flight.Client and return a flightsql.Client containing it rather // than having to manually construct both yourself. It just delegates @@ -47,6 +51,8 @@ func NewClient(addr string, auth flight.ClientAuthHandler, middleware []flight.C // Client wraps a regular Flight RPC Client to provide the FlightSQL // interface functions and methods. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) type Client struct { Client flight.Client } @@ -72,7 +78,11 @@ func flightInfoForCommand(ctx context.Context, cl *Client, cmd proto.Message, op if err != nil { return nil, err } +<<<<<<< HEAD return cl.getFlightInfo(ctx, desc, opts...) +======= + return cl.GetFlightInfo(ctx, desc, opts...) +>>>>>>> 7c0a7b52f (initial flightsql client) } // Execute executes the desired query on the server and returns a FlightInfo @@ -85,7 +95,11 @@ func (c *Client) Execute(ctx context.Context, query string, opts ...grpc.CallOpt // ExecuteUpdate is for executing an update query and only returns the number of affected rows. func (c *Client) ExecuteUpdate(ctx context.Context, query string, opts ...grpc.CallOption) (n int64, err error) { var ( +<<<<<<< HEAD cmd pb.CommandStatementUpdate +======= + cmd pb.CommandStatementQuery +>>>>>>> 7c0a7b52f (initial flightsql client) desc *flight.FlightDescriptor stream pb.FlightService_DoPutClient res *pb.PutResult @@ -129,7 +143,11 @@ func (c *Client) GetCatalogs(ctx context.Context, opts ...grpc.CallOption) (*fli // GetDBSchemas requests the list of schemas from the database and // returns a FlightInfo object where the response can be retrieved func (c *Client) GetDBSchemas(ctx context.Context, cmdOpts *GetDBSchemasOpts, opts ...grpc.CallOption) (*flight.FlightInfo, error) { +<<<<<<< HEAD return flightInfoForCommand(ctx, c, (*pb.CommandGetDbSchemas)(cmdOpts), opts...) +======= + return flightInfoForCommand(ctx, c, cmdOpts, opts...) +>>>>>>> 7c0a7b52f (initial flightsql client) } // DoGet uses the provided flight ticket to request the stream of data. @@ -149,12 +167,18 @@ func (c *Client) DoGet(ctx context.Context, in *flight.Ticket, opts ...grpc.Call // should be returned, etc.). Returns a FlightInfo object where the response // can be retrieved. func (c *Client) GetTables(ctx context.Context, reqOptions *GetTablesOpts, opts ...grpc.CallOption) (*flight.FlightInfo, error) { +<<<<<<< HEAD return flightInfoForCommand(ctx, c, (*pb.CommandGetTables)(reqOptions), opts...) } // GetPrimaryKeys requests the primary keys for a specific table from the // server, specified using a TableRef. Returns a FlightInfo object where // the response can be retrieved. +======= + return flightInfoForCommand(ctx, c, reqOptions, opts...) +} + +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetPrimaryKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetPrimaryKeys{ Catalog: ref.Catalog, @@ -164,9 +188,12 @@ func (c *Client) GetPrimaryKeys(ctx context.Context, ref TableRef, opts ...grpc. return flightInfoForCommand(ctx, c, &cmd, opts...) } +<<<<<<< HEAD // GetExportedKeys retrieves a description about the foreign key columns // that reference the primary key columns of the specified table. Returns // a FlightInfo object where the response can be retrieved. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetExportedKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetExportedKeys{ Catalog: ref.Catalog, @@ -176,8 +203,11 @@ func (c *Client) GetExportedKeys(ctx context.Context, ref TableRef, opts ...grpc return flightInfoForCommand(ctx, c, &cmd, opts...) } +<<<<<<< HEAD // GetImportedKeys returns the foreign key columns for the specified table. // Returns a FlightInfo object indicating where the response can be retrieved. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetImportedKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetImportedKeys{ Catalog: ref.Catalog, @@ -187,11 +217,14 @@ func (c *Client) GetImportedKeys(ctx context.Context, ref TableRef, opts ...grpc return flightInfoForCommand(ctx, c, &cmd, opts...) } +<<<<<<< HEAD // GetCrossReference retrieves a description of the foreign key columns // in the specified ForeignKey table that reference the primary key or // columns representing a restraint of the parent table (could be the same // or a different table). Returns a FlightInfo object indicating where // the response can be retrieved with DoGet. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetCrossReference(ctx context.Context, pkTable, fkTable TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetCrossReference{ PkCatalog: pkTable.Catalog, @@ -204,23 +237,32 @@ func (c *Client) GetCrossReference(ctx context.Context, pkTable, fkTable TableRe return flightInfoForCommand(ctx, c, &cmd, opts...) } +<<<<<<< HEAD // GetTableTypes requests a list of the types of tables available on this // server. Returns a FlightInfo object indicating where the response can // be retrieved. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetTableTypes(ctx context.Context, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return flightInfoForCommand(ctx, c, &pb.CommandGetTableTypes{}, opts...) } +<<<<<<< HEAD // GetXdbcTypeInfo requests the information about all the data types supported // (dataType == nil) or a specific data type. Returns a FlightInfo object // indicating where the response can be retrieved. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetXdbcTypeInfo(ctx context.Context, dataType *int32, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return flightInfoForCommand(ctx, c, &pb.CommandGetXdbcTypeInfo{DataType: dataType}, opts...) } +<<<<<<< HEAD // GetSqlInfo returns a list of the requested SQL information corresponding // to the values in the info slice. Returns a FlightInfo object indicating // where the response can be retrieved. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := &pb.CommandGetSqlInfo{Info: make([]uint32, len(info))} @@ -230,12 +272,17 @@ func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.Ca return flightInfoForCommand(ctx, c, cmd, opts...) } +<<<<<<< HEAD // Prepare creates a PreparedStatement object for the specified query. // The resulting PreparedStatement object should be Closed when no longer // needed. It will maintain a reference to this Client for use to execute // and use the specified allocator for any allocations it needs to perform. func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string, opts ...grpc.CallOption) (prep *PreparedStatement, err error) { const actionType = CreatePreparedStatementActionType +======= +func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string, opts ...grpc.CallOption) (prep *PreparedStatement, err error) { + const actionType = "CreatePreparedStatement" +>>>>>>> 7c0a7b52f (initial flightsql client) var ( cmd, cmdResult anypb.Any res *pb.Result @@ -272,6 +319,7 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string return } +<<<<<<< HEAD if result.DatasetSchema != nil { dsSchema, err = flight.DeserializeSchema(result.DatasetSchema, mem) if err != nil { @@ -283,6 +331,15 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string if err != nil { return } +======= + dsSchema, err = flight.DeserializeSchema(result.DatasetSchema, mem) + if err != nil { + return + } + paramSchema, err = flight.DeserializeSchema(result.ParameterSchema, mem) + if err != nil { + return +>>>>>>> 7c0a7b52f (initial flightsql client) } prep = &PreparedStatement{ @@ -295,6 +352,7 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string return } +<<<<<<< HEAD func (c *Client) getFlightInfo(ctx context.Context, desc *flight.FlightDescriptor, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return c.Client.GetFlightInfo(ctx, desc, opts...) } @@ -309,6 +367,14 @@ func (c *Client) Close() error { return c.Client.Close() } // If the server returned the Dataset Schema or Parameter Binding schemas // at creation, they will also be accessible from this object. Close // should be called when no longer needed. +======= +func (c *Client) GetFlightInfo(ctx context.Context, desc *flight.FlightDescriptor, opts ...grpc.CallOption) (*flight.FlightInfo, error) { + return c.Client.GetFlightInfo(ctx, desc, opts...) +} + +func (c *Client) Close() error { return c.Client.Close() } + +>>>>>>> 7c0a7b52f (initial flightsql client) type PreparedStatement struct { client *Client opts []grpc.CallOption @@ -319,11 +385,14 @@ type PreparedStatement struct { closed bool } +<<<<<<< HEAD // Execute executes the prepared statement on the server and returns a FlightInfo // indicating where to retrieve the response. If SetParameters has been called // then the parameter bindings will be sent before execution. // // Will error if already closed. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) Execute(ctx context.Context) (*flight.FlightInfo, error) { if p.closed { return nil, errors.New("arrow/flightsql: prepared statement already closed") @@ -353,17 +422,27 @@ func (p *PreparedStatement) Execute(ctx context.Context) (*flight.FlightInfo, er pstream.CloseSend() // wait for the server to ack the result +<<<<<<< HEAD if _, err = pstream.Recv(); err != nil && err != io.EOF { +======= + if _, err = pstream.Recv(); err != nil { +>>>>>>> 7c0a7b52f (initial flightsql client) return nil, err } } +<<<<<<< HEAD return p.client.getFlightInfo(ctx, desc, p.opts...) } // ExecuteUpdate executes the prepared statement update query on the server // and returns the number of rows affected. If SetParameters was called, // the parameter bindings will be sent with the request to execute. +======= + return p.client.GetFlightInfo(ctx, desc, p.opts...) +} + +>>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, err error) { if p.closed { return 0, errors.New("arrow/flightsql: prepared statement already closed") @@ -394,7 +473,11 @@ func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, } } else { schema := arrow.NewSchema([]arrow.Field{}, nil) +<<<<<<< HEAD wr = flight.NewRecordWriter(pstream, ipc.WithSchema(schema)) +======= + wr = flight.NewRecordWriter(pstream, ipc.WithSchema(p.paramBinding.Schema())) +>>>>>>> 7c0a7b52f (initial flightsql client) wr.SetFlightDescriptor(desc) rec := array.NewRecord(schema, []arrow.Array{}, 0) if err = wr.Write(rec); err != nil { @@ -419,6 +502,7 @@ func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, return updateResult.GetRecordCount(), nil } +<<<<<<< HEAD // DatasetSchema may be nil if the server did not return it when creating the // Prepared Statement. func (p *PreparedStatement) DatasetSchema() *arrow.Schema { return p.datasetSchema } @@ -434,6 +518,11 @@ func (p *PreparedStatement) ParameterSchema() *arrow.Schema { return p.paramSche // out from under the statement. Release will be called on a previous // binding record if it existed, and will be called upon calling Close // on the PreparedStatement. +======= +func (p *PreparedStatement) DatasetSchema() *arrow.Schema { return p.datasetSchema } +func (p *PreparedStatement) ParameterSchema() *arrow.Schema { return p.paramSchema } + +>>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) SetParameters(binding arrow.Record) { if p.paramBinding != nil { p.paramBinding.Release() @@ -443,9 +532,12 @@ func (p *PreparedStatement) SetParameters(binding arrow.Record) { p.paramBinding.Retain() } +<<<<<<< HEAD // Close calls release on any parameter binding record and sends // a ClosePreparedStatement action to the server. After calling // Close, the PreparedStatement should not be used again. +======= +>>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) Close(ctx context.Context) error { if p.closed { return errors.New("arrow/flightsql: already closed") @@ -456,7 +548,11 @@ func (p *PreparedStatement) Close(ctx context.Context) error { p.paramBinding = nil } +<<<<<<< HEAD const actionType = ClosePreparedStatementActionType +======= + const actionType = "ClosePreparedStatement" +>>>>>>> 7c0a7b52f (initial flightsql client) var ( cmd anypb.Any request pb.ActionClosePreparedStatementRequest diff --git a/go/arrow/flight/flightsql/types.go b/go/arrow/flight/flightsql/types.go index 5e033d00ee322..956775325e68b 100644 --- a/go/arrow/flight/flightsql/types.go +++ b/go/arrow/flight/flightsql/types.go @@ -18,114 +18,31 @@ package flightsql import ( pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" - "google.golang.org/protobuf/proto" - "google.golang.org/protobuf/types/known/anypb" ) -// Constants for Action types -const ( - CreatePreparedStatementActionType = "CreatePreparedStatement" - ClosePreparedStatementActionType = "ClosePreparedStatement" -) - -func toCrossTableRef(cmd *pb.CommandGetCrossReference) CrossTableRef { - return CrossTableRef{ - PKRef: TableRef{ - Catalog: cmd.PkCatalog, - DBSchema: cmd.PkDbSchema, - Table: cmd.PkTable, - }, - FKRef: TableRef{ - Catalog: cmd.FkCatalog, - DBSchema: cmd.FkDbSchema, - Table: cmd.FkTable, - }, - } -} - -func pkToTableRef(cmd *pb.CommandGetPrimaryKeys) TableRef { - return TableRef{ - Catalog: cmd.Catalog, - DBSchema: cmd.DbSchema, - Table: cmd.Table, - } -} - -func exkToTableRef(cmd *pb.CommandGetExportedKeys) TableRef { - return TableRef{ - Catalog: cmd.Catalog, - DBSchema: cmd.DbSchema, - Table: cmd.Table, - } -} - -func impkToTableRef(cmd *pb.CommandGetImportedKeys) TableRef { - return TableRef{ - Catalog: cmd.Catalog, - DBSchema: cmd.DbSchema, - Table: cmd.Table, - } -} - -// CreateStatementQueryTicket is a helper that constructs a properly -// serialized TicketStatementQuery containing a given opaque binary handle -// for use with constructing a ticket to return from GetFlightInfoStatement. -func CreateStatementQueryTicket(handle []byte) ([]byte, error) { - query := &pb.TicketStatementQuery{StatementHandle: handle} - var ticket anypb.Any - ticket.MarshalFrom(query) - - return proto.Marshal(&ticket) +type TableRef struct { + // Catalog specifies the catalog this table belongs to. + // An empty string refers to tables without a catalog. + // If nil, can reference a table in any catalog. + Catalog *string + // DBSchema specifies the database schema the table belongs to. + // An empty string refers to a table which does not belong to + // a database schema. + // If nil, can reference a table in any database schema. + DBSchema *string + // Table is the name of the table that is being referenced. + Table string } type ( - // GetDBSchemasOpts contains the options to request Database Schemas: - // an optional Catalog and a Schema Name filter pattern. - GetDBSchemasOpts pb.CommandGetDbSchemas - // GetTablesOpts contains the options for retrieving a list of tables: - // optional Catalog, Schema filter pattern, Table name filter pattern, - // a filter of table types, and whether or not to include the schema - // in the response. - GetTablesOpts pb.CommandGetTables - - // SqlInfoResultMap is a mapping of SqlInfo ids to the desired response. - // This is part of a Server and used for registering responses to a - // SqlInfo request. - SqlInfoResultMap map[uint32]interface{} - - // TableRef is a helpful struct for referencing a specific Table - // by its catalog, schema, and table name. - TableRef struct { - // Catalog specifies the catalog this table belongs to. - // An empty string refers to tables without a catalog. - // If nil, can reference a table in any catalog. - Catalog *string - // DBSchema specifies the database schema the table belongs to. - // An empty string refers to a table which does not belong to - // a database schema. - // If nil, can reference a table in any database schema. - DBSchema *string - // Table is the name of the table that is being referenced. - Table string - } - - // CrossTableRef contains a reference to a Primary Key table - // and a Foreign Key table. - CrossTableRef struct { - PKRef TableRef - FKRef TableRef - } - + GetDBSchemasOpts = pb.CommandGetDbSchemas + GetTablesOpts = pb.CommandGetTables // since we are hiding the Protobuf internals in an internal // package, we need to provide enum values for the SqlInfo enum here SqlInfo uint32 ) -// SqlInfo enum values const ( - // Server Information - // Values [0-500): Provide information about the Flight SQL Server itself - // Retrieves a UTF-8 string with the name of the Flight SQL Server. SqlInfoFlightSqlServerName = SqlInfo(pb.SqlInfo_FLIGHT_SQL_SERVER_NAME) // Retrieves a UTF-8 string with the native version of the Flight SQL Server. @@ -140,9 +57,6 @@ const ( // - true: if read only SqlInfoFlightSqlServerReadOnly = SqlInfo(pb.SqlInfo_FLIGHT_SQL_SERVER_READ_ONLY) - // SQL Syntax Information - // Values [500-1000): provide information about the supported SQL Syntax - // Retrieves a boolean value indicating whether the Flight SQL Server supports CREATE and DROP of catalogs. // // Returns: @@ -686,60 +600,3 @@ const ( ) func (s SqlInfo) String() string { return pb.SqlInfo(int32(s)).String() } - -// SqlSupportedCaseSensitivity indicates whether something -// (e.g. an identifier) is case-sensitive -// -// duplicated from protobuf to avoid relying directly on the protobuf -// generated code, also making them shorter and easier to use -type SqlSupportedCaseSensitivity = pb.SqlSupportedCaseSensitivity - -const ( - SqlCaseSensitivityUnknown = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_UNKNOWN - SqlCaseSensitivityCaseInsensitive = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_CASE_INSENSITIVE - SqlCaseSensitivityUpperCase = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_UPPERCASE - SqlCaseSensitivityLowerCase = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_LOWERCASE -) - -// SqlNullOrdering indicates how nulls are sorted -// -// duplicated from protobuf to avoid relying directly on the protobuf -// generated code, also making them shorter and easier to use -type SqlNullOrdering = pb.SqlNullOrdering - -const ( - SqlNullOrderingSortHigh = pb.SqlNullOrdering_SQL_NULLS_SORTED_HIGH - SqlNullOrderingSortLow = pb.SqlNullOrdering_SQL_NULLS_SORTED_LOW - SqlNullOrderingSortAtStart = pb.SqlNullOrdering_SQL_NULLS_SORTED_AT_START - SqlNullOrderingSortAtEnd = pb.SqlNullOrdering_SQL_NULLS_SORTED_AT_END -) - -// SqlSupportsConvert indicates support for converting between different -// types. -// -// duplicated from protobuf to avoid relying directly on the protobuf -// generated code, also making them shorter and easier to use -type SqlSupportsConvert = pb.SqlSupportsConvert - -const ( - SqlConvertBigInt = pb.SqlSupportsConvert_SQL_CONVERT_BIGINT - SqlConvertBinary = pb.SqlSupportsConvert_SQL_CONVERT_BINARY - SqlConvertBit = pb.SqlSupportsConvert_SQL_CONVERT_BIT - SqlConvertChar = pb.SqlSupportsConvert_SQL_CONVERT_CHAR - SqlConvertDate = pb.SqlSupportsConvert_SQL_CONVERT_DATE - SqlConvertDecimal = pb.SqlSupportsConvert_SQL_CONVERT_DECIMAL - SqlConvertFloat = pb.SqlSupportsConvert_SQL_CONVERT_FLOAT - SqlConvertInteger = pb.SqlSupportsConvert_SQL_CONVERT_INTEGER - SqlConvertIntervalDayTime = pb.SqlSupportsConvert_SQL_CONVERT_INTERVAL_DAY_TIME - SqlConvertIntervalYearMonth = pb.SqlSupportsConvert_SQL_CONVERT_INTERVAL_YEAR_MONTH - SqlConvertLongVarbinary = pb.SqlSupportsConvert_SQL_CONVERT_LONGVARBINARY - SqlConvertLongVarchar = pb.SqlSupportsConvert_SQL_CONVERT_LONGVARCHAR - SqlConvertNumeric = pb.SqlSupportsConvert_SQL_CONVERT_NUMERIC - SqlConvertReal = pb.SqlSupportsConvert_SQL_CONVERT_REAL - SqlConvertSmallInt = pb.SqlSupportsConvert_SQL_CONVERT_SMALLINT - SqlConvertTime = pb.SqlSupportsConvert_SQL_CONVERT_TIME - SqlConvertTimestamp = pb.SqlSupportsConvert_SQL_CONVERT_TIMESTAMP - SqlConvertTinyInt = pb.SqlSupportsConvert_SQL_CONVERT_TINYINT - SqlConvertVarbinary = pb.SqlSupportsConvert_SQL_CONVERT_VARBINARY - SqlConvertVarchar = pb.SqlSupportsConvert_SQL_CONVERT_VARCHAR -) From 3384cafd7815832d1cce852e6e1a9b7f9fbc665f Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 8 Aug 2022 16:45:29 -0400 Subject: [PATCH 04/23] initial implementation --- go/arrow/flight/flightsql/client.go | 10 +- .../flightsql/schema_ref/reference_schemas.go | 6 +- go/arrow/flight/flightsql/types.go | 140 ++++++++++++++++-- 3 files changed, 132 insertions(+), 24 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index b4a9bbf970e72..79eca52399462 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -272,17 +272,13 @@ func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.Ca return flightInfoForCommand(ctx, c, cmd, opts...) } -<<<<<<< HEAD // Prepare creates a PreparedStatement object for the specified query. // The resulting PreparedStatement object should be Closed when no longer // needed. It will maintain a reference to this Client for use to execute // and use the specified allocator for any allocations it needs to perform. func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string, opts ...grpc.CallOption) (prep *PreparedStatement, err error) { const actionType = CreatePreparedStatementActionType -======= -func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string, opts ...grpc.CallOption) (prep *PreparedStatement, err error) { - const actionType = "CreatePreparedStatement" ->>>>>>> 7c0a7b52f (initial flightsql client) + var ( cmd, cmdResult anypb.Any res *pb.Result @@ -548,11 +544,7 @@ func (p *PreparedStatement) Close(ctx context.Context) error { p.paramBinding = nil } -<<<<<<< HEAD const actionType = ClosePreparedStatementActionType -======= - const actionType = "ClosePreparedStatement" ->>>>>>> 7c0a7b52f (initial flightsql client) var ( cmd anypb.Any request pb.ActionClosePreparedStatementRequest diff --git a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go index 7a4a14064d540..f2240d4500b00 100644 --- a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go +++ b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go @@ -64,7 +64,7 @@ var ( {Name: "fk_key_name", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "pk_key_name", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "update_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, - {Name: "delete_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, + {Name: "delete_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: true}, }, nil) ImportedKeys = ImportedExportedKeysAndCrossReference ExportedKeys = ImportedExportedKeysAndCrossReference @@ -80,7 +80,7 @@ var ( {Name: "int32_to_int32_list_map", Type: arrow.MapOf(arrow.PrimitiveTypes.Int32, arrow.ListOf(arrow.PrimitiveTypes.Int32))}, - }, []arrow.UnionTypeCode{0, 1, 2, 3, 4, 5})}, + }, []arrow.UnionTypeCode{})}, }, nil) XdbcTypeInfo = arrow.NewSchema([]arrow.Field{ {Name: "type_name", Type: arrow.BinaryTypes.String, Nullable: false}, @@ -88,7 +88,7 @@ var ( {Name: "column_size", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "literal_prefix", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "literal_suffix", Type: arrow.BinaryTypes.String, Nullable: true}, - {Name: "create_params", Type: arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.BinaryTypes.String, Nullable: false}), Nullable: true}, + {Name: "create_params", Type: arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.BinaryTypes.String, Nullable: false})}, {Name: "nullable", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, {Name: "case_sensitive", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, {Name: "searchable", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, diff --git a/go/arrow/flight/flightsql/types.go b/go/arrow/flight/flightsql/types.go index 956775325e68b..6f11395a7facd 100644 --- a/go/arrow/flight/flightsql/types.go +++ b/go/arrow/flight/flightsql/types.go @@ -20,29 +20,85 @@ import ( pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" ) -type TableRef struct { - // Catalog specifies the catalog this table belongs to. - // An empty string refers to tables without a catalog. - // If nil, can reference a table in any catalog. - Catalog *string - // DBSchema specifies the database schema the table belongs to. - // An empty string refers to a table which does not belong to - // a database schema. - // If nil, can reference a table in any database schema. - DBSchema *string - // Table is the name of the table that is being referenced. - Table string +const ( + CreatePreparedStatementActionType = "CreatePreparedStatement" + ClosePreparedStatementActionType = "ClosePreparedStatement" +) + +func toCrossTableRef(cmd *pb.CommandGetCrossReference) CrossTableRef { + return CrossTableRef{ + PKRef: TableRef{ + Catalog: cmd.PkCatalog, + DBSchema: cmd.PkDbSchema, + Table: cmd.PkTable, + }, + FKRef: TableRef{ + Catalog: cmd.FkCatalog, + DBSchema: cmd.FkDbSchema, + Table: cmd.FkTable, + }, + } +} + +func pkToTableRef(cmd *pb.CommandGetPrimaryKeys) TableRef { + return TableRef{ + Catalog: cmd.Catalog, + DBSchema: cmd.DbSchema, + Table: cmd.Table, + } +} + +func exkToTableRef(cmd *pb.CommandGetExportedKeys) TableRef { + return TableRef{ + Catalog: cmd.Catalog, + DBSchema: cmd.DbSchema, + Table: cmd.Table, + } +} + +func impkToTableRef(cmd *pb.CommandGetImportedKeys) TableRef { + return TableRef{ + Catalog: cmd.Catalog, + DBSchema: cmd.DbSchema, + Table: cmd.Table, + } } type ( GetDBSchemasOpts = pb.CommandGetDbSchemas GetTablesOpts = pb.CommandGetTables + + SqlInfoResultMap map[uint32]interface{} + + TableRef struct { + // Catalog specifies the catalog this table belongs to. + // An empty string refers to tables without a catalog. + // If nil, can reference a table in any catalog. + Catalog *string + // DBSchema specifies the database schema the table belongs to. + // An empty string refers to a table which does not belong to + // a database schema. + // If nil, can reference a table in any database schema. + DBSchema *string + // Table is the name of the table that is being referenced. + Table string + } + + CrossTableRef struct { + PKRef TableRef + FKRef TableRef + } + // since we are hiding the Protobuf internals in an internal // package, we need to provide enum values for the SqlInfo enum here SqlInfo uint32 ) +// SqlInfo enum values const ( + // Server Information + // Values [0-500): Provide information about the Flight SQL Server itself + // Retrieves a UTF-8 string with the name of the Flight SQL Server. SqlInfoFlightSqlServerName = SqlInfo(pb.SqlInfo_FLIGHT_SQL_SERVER_NAME) // Retrieves a UTF-8 string with the native version of the Flight SQL Server. @@ -57,6 +113,9 @@ const ( // - true: if read only SqlInfoFlightSqlServerReadOnly = SqlInfo(pb.SqlInfo_FLIGHT_SQL_SERVER_READ_ONLY) + // SQL Syntax Information + // Values [500-1000): provide information about the supported SQL Syntax + // Retrieves a boolean value indicating whether the Flight SQL Server supports CREATE and DROP of catalogs. // // Returns: @@ -600,3 +659,60 @@ const ( ) func (s SqlInfo) String() string { return pb.SqlInfo(int32(s)).String() } + +// SqlSupportedCaseSensitivity indicates whether something +// (e.g. an identifier) is case-sensitive +// +// duplicated from protobuf to avoid relying directly on the protobuf +// generated code, also making them shorter and easier to use +type SqlSupportedCaseSensitivity = pb.SqlSupportedCaseSensitivity + +const ( + SqlCaseSensitivityUnknown = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_UNKNOWN + SqlCaseSensitivityCaseInsensitive = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_CASE_INSENSITIVE + SqlCaseSensitivityUpperCase = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_UPPERCASE + SqlCaseSensitivityLowerCase = pb.SqlSupportedCaseSensitivity_SQL_CASE_SENSITIVITY_LOWERCASE +) + +// SqlNullOrdering indicates how nulls are sorted +// +// duplicated from protobuf to avoid relying directly on the protobuf +// generated code, also making them shorter and easier to use +type SqlNullOrdering = pb.SqlNullOrdering + +const ( + SqlNullOrderingSortHigh = pb.SqlNullOrdering_SQL_NULLS_SORTED_HIGH + SqlNullOrderingSortLow = pb.SqlNullOrdering_SQL_NULLS_SORTED_LOW + SqlNullOrderingSortAtStart = pb.SqlNullOrdering_SQL_NULLS_SORTED_AT_START + SqlNullOrderingSortAtEnd = pb.SqlNullOrdering_SQL_NULLS_SORTED_AT_END +) + +// SqlSupportsConvert indicates support for converting between different +// types. +// +// duplicated from protobuf to avoid relying directly on the protobuf +// generated code, also making them shorter and easier to use +type SqlSupportsConvert = pb.SqlSupportsConvert + +const ( + SqlConvertBigInt = pb.SqlSupportsConvert_SQL_CONVERT_BIGINT + SqlConvertBinary = pb.SqlSupportsConvert_SQL_CONVERT_BINARY + SqlConvertBit = pb.SqlSupportsConvert_SQL_CONVERT_BIT + SqlConvertChar = pb.SqlSupportsConvert_SQL_CONVERT_CHAR + SqlConvertDate = pb.SqlSupportsConvert_SQL_CONVERT_DATE + SqlConvertDecimal = pb.SqlSupportsConvert_SQL_CONVERT_DECIMAL + SqlConvertFloat = pb.SqlSupportsConvert_SQL_CONVERT_FLOAT + SqlConvertInteger = pb.SqlSupportsConvert_SQL_CONVERT_INTEGER + SqlConvertIntervalDayTime = pb.SqlSupportsConvert_SQL_CONVERT_INTERVAL_DAY_TIME + SqlConvertIntervalYearMonth = pb.SqlSupportsConvert_SQL_CONVERT_INTERVAL_YEAR_MONTH + SqlConvertLongVarbinary = pb.SqlSupportsConvert_SQL_CONVERT_LONGVARBINARY + SqlConvertLongVarchar = pb.SqlSupportsConvert_SQL_CONVERT_LONGVARCHAR + SqlConvertNumeric = pb.SqlSupportsConvert_SQL_CONVERT_NUMERIC + SqlConvertReal = pb.SqlSupportsConvert_SQL_CONVERT_REAL + SqlConvertSmallInt = pb.SqlSupportsConvert_SQL_CONVERT_SMALLINT + SqlConvertTime = pb.SqlSupportsConvert_SQL_CONVERT_TIME + SqlConvertTimestamp = pb.SqlSupportsConvert_SQL_CONVERT_TIMESTAMP + SqlConvertTinyInt = pb.SqlSupportsConvert_SQL_CONVERT_TINYINT + SqlConvertVarbinary = pb.SqlSupportsConvert_SQL_CONVERT_VARBINARY + SqlConvertVarchar = pb.SqlSupportsConvert_SQL_CONVERT_VARCHAR +) From 2f2b18adebbcbd520b611ea208df0daba6a3cd62 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 8 Aug 2022 19:10:37 -0400 Subject: [PATCH 05/23] client unit tests --- go/arrow/flight/flightsql/client.go | 14 -------------- .../flightsql/schema_ref/reference_schemas.go | 2 +- 2 files changed, 1 insertion(+), 15 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index 79eca52399462..9eecf6963e90e 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -95,11 +95,7 @@ func (c *Client) Execute(ctx context.Context, query string, opts ...grpc.CallOpt // ExecuteUpdate is for executing an update query and only returns the number of affected rows. func (c *Client) ExecuteUpdate(ctx context.Context, query string, opts ...grpc.CallOption) (n int64, err error) { var ( -<<<<<<< HEAD cmd pb.CommandStatementUpdate -======= - cmd pb.CommandStatementQuery ->>>>>>> 7c0a7b52f (initial flightsql client) desc *flight.FlightDescriptor stream pb.FlightService_DoPutClient res *pb.PutResult @@ -315,7 +311,6 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string return } -<<<<<<< HEAD if result.DatasetSchema != nil { dsSchema, err = flight.DeserializeSchema(result.DatasetSchema, mem) if err != nil { @@ -327,15 +322,6 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string if err != nil { return } -======= - dsSchema, err = flight.DeserializeSchema(result.DatasetSchema, mem) - if err != nil { - return - } - paramSchema, err = flight.DeserializeSchema(result.ParameterSchema, mem) - if err != nil { - return ->>>>>>> 7c0a7b52f (initial flightsql client) } prep = &PreparedStatement{ diff --git a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go index f2240d4500b00..728280b5dee01 100644 --- a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go +++ b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go @@ -80,7 +80,7 @@ var ( {Name: "int32_to_int32_list_map", Type: arrow.MapOf(arrow.PrimitiveTypes.Int32, arrow.ListOf(arrow.PrimitiveTypes.Int32))}, - }, []arrow.UnionTypeCode{})}, + }, []arrow.UnionTypeCode{0, 1, 2, 3, 4, 5})}, }, nil) XdbcTypeInfo = arrow.NewSchema([]arrow.Field{ {Name: "type_name", Type: arrow.BinaryTypes.String, Nullable: false}, From c83e611d37913cbf197975c20bbd85f173ea0079 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 8 Aug 2022 19:38:57 -0400 Subject: [PATCH 06/23] basic test framework for server --- go/arrow/flight/flightsql/client.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index 9eecf6963e90e..ff94ec6382ef1 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -35,7 +35,6 @@ import ( "google.golang.org/protobuf/types/known/anypb" ) -<<<<<<< HEAD // NewClient is a convenience function to automatically construct // a flight.Client and return a flightsql.Client containing it rather // than having to manually construct both yourself. It just delegates @@ -51,8 +50,6 @@ func NewClient(addr string, auth flight.ClientAuthHandler, middleware []flight.C // Client wraps a regular Flight RPC Client to provide the FlightSQL // interface functions and methods. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) type Client struct { Client flight.Client } From 88cc0a21ed9621e414cb7f94ef293feb1a479045 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Tue, 9 Aug 2022 12:35:02 -0400 Subject: [PATCH 07/23] get flight_sql integration up and running --- go/arrow/array/record_test.go | 2 +- go/arrow/flight/flightsql/client.go | 17 ----------------- go/arrow/flight/flightsql/types.go | 10 ++++++++++ 3 files changed, 11 insertions(+), 18 deletions(-) diff --git a/go/arrow/array/record_test.go b/go/arrow/array/record_test.go index d5fbeb8c8e6dc..5deeb27853b73 100644 --- a/go/arrow/array/record_test.go +++ b/go/arrow/array/record_test.go @@ -135,7 +135,7 @@ func TestRecord(t *testing.T) { { schema: schema, cols: nil, - rows: 0, + rows: 0, }, { schema: schema, diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index ff94ec6382ef1..bf9e9558f8fb9 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -19,10 +19,7 @@ package flightsql import ( "context" "errors" -<<<<<<< HEAD "io" -======= ->>>>>>> 7c0a7b52f (initial flightsql client) "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/array" @@ -401,27 +398,17 @@ func (p *PreparedStatement) Execute(ctx context.Context) (*flight.FlightInfo, er pstream.CloseSend() // wait for the server to ack the result -<<<<<<< HEAD if _, err = pstream.Recv(); err != nil && err != io.EOF { -======= - if _, err = pstream.Recv(); err != nil { ->>>>>>> 7c0a7b52f (initial flightsql client) return nil, err } } -<<<<<<< HEAD return p.client.getFlightInfo(ctx, desc, p.opts...) } // ExecuteUpdate executes the prepared statement update query on the server // and returns the number of rows affected. If SetParameters was called, // the parameter bindings will be sent with the request to execute. -======= - return p.client.GetFlightInfo(ctx, desc, p.opts...) -} - ->>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, err error) { if p.closed { return 0, errors.New("arrow/flightsql: prepared statement already closed") @@ -452,11 +439,7 @@ func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, } } else { schema := arrow.NewSchema([]arrow.Field{}, nil) -<<<<<<< HEAD wr = flight.NewRecordWriter(pstream, ipc.WithSchema(schema)) -======= - wr = flight.NewRecordWriter(pstream, ipc.WithSchema(p.paramBinding.Schema())) ->>>>>>> 7c0a7b52f (initial flightsql client) wr.SetFlightDescriptor(desc) rec := array.NewRecord(schema, []arrow.Array{}, 0) if err = wr.Write(rec); err != nil { diff --git a/go/arrow/flight/flightsql/types.go b/go/arrow/flight/flightsql/types.go index 6f11395a7facd..96bcf8073ef15 100644 --- a/go/arrow/flight/flightsql/types.go +++ b/go/arrow/flight/flightsql/types.go @@ -18,6 +18,8 @@ package flightsql import ( pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/anypb" ) const ( @@ -64,6 +66,14 @@ func impkToTableRef(cmd *pb.CommandGetImportedKeys) TableRef { } } +func CreateStatementQueryTicket(handle []byte) ([]byte, error) { + query := &pb.TicketStatementQuery{StatementHandle: handle} + var ticket anypb.Any + ticket.MarshalFrom(query) + + return proto.Marshal(&ticket) +} + type ( GetDBSchemasOpts = pb.CommandGetDbSchemas GetTablesOpts = pb.CommandGetTables From 936dfd1e9ddeeba23c5f362f22aa99207b3b2ceb Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Tue, 9 Aug 2022 14:04:17 -0400 Subject: [PATCH 08/23] add lots of comments --- go/arrow/flight/flightsql/client.go | 57 ++--------------------------- go/arrow/flight/flightsql/types.go | 23 ++++++++++-- 2 files changed, 24 insertions(+), 56 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index bf9e9558f8fb9..080cd039a71c7 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -72,11 +72,7 @@ func flightInfoForCommand(ctx context.Context, cl *Client, cmd proto.Message, op if err != nil { return nil, err } -<<<<<<< HEAD return cl.getFlightInfo(ctx, desc, opts...) -======= - return cl.GetFlightInfo(ctx, desc, opts...) ->>>>>>> 7c0a7b52f (initial flightsql client) } // Execute executes the desired query on the server and returns a FlightInfo @@ -133,11 +129,7 @@ func (c *Client) GetCatalogs(ctx context.Context, opts ...grpc.CallOption) (*fli // GetDBSchemas requests the list of schemas from the database and // returns a FlightInfo object where the response can be retrieved func (c *Client) GetDBSchemas(ctx context.Context, cmdOpts *GetDBSchemasOpts, opts ...grpc.CallOption) (*flight.FlightInfo, error) { -<<<<<<< HEAD return flightInfoForCommand(ctx, c, (*pb.CommandGetDbSchemas)(cmdOpts), opts...) -======= - return flightInfoForCommand(ctx, c, cmdOpts, opts...) ->>>>>>> 7c0a7b52f (initial flightsql client) } // DoGet uses the provided flight ticket to request the stream of data. @@ -157,18 +149,12 @@ func (c *Client) DoGet(ctx context.Context, in *flight.Ticket, opts ...grpc.Call // should be returned, etc.). Returns a FlightInfo object where the response // can be retrieved. func (c *Client) GetTables(ctx context.Context, reqOptions *GetTablesOpts, opts ...grpc.CallOption) (*flight.FlightInfo, error) { -<<<<<<< HEAD return flightInfoForCommand(ctx, c, (*pb.CommandGetTables)(reqOptions), opts...) } // GetPrimaryKeys requests the primary keys for a specific table from the // server, specified using a TableRef. Returns a FlightInfo object where // the response can be retrieved. -======= - return flightInfoForCommand(ctx, c, reqOptions, opts...) -} - ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetPrimaryKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetPrimaryKeys{ Catalog: ref.Catalog, @@ -178,12 +164,9 @@ func (c *Client) GetPrimaryKeys(ctx context.Context, ref TableRef, opts ...grpc. return flightInfoForCommand(ctx, c, &cmd, opts...) } -<<<<<<< HEAD // GetExportedKeys retrieves a description about the foreign key columns // that reference the primary key columns of the specified table. Returns // a FlightInfo object where the response can be retrieved. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetExportedKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetExportedKeys{ Catalog: ref.Catalog, @@ -193,11 +176,8 @@ func (c *Client) GetExportedKeys(ctx context.Context, ref TableRef, opts ...grpc return flightInfoForCommand(ctx, c, &cmd, opts...) } -<<<<<<< HEAD // GetImportedKeys returns the foreign key columns for the specified table. // Returns a FlightInfo object indicating where the response can be retrieved. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetImportedKeys(ctx context.Context, ref TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetImportedKeys{ Catalog: ref.Catalog, @@ -207,14 +187,11 @@ func (c *Client) GetImportedKeys(ctx context.Context, ref TableRef, opts ...grpc return flightInfoForCommand(ctx, c, &cmd, opts...) } -<<<<<<< HEAD // GetCrossReference retrieves a description of the foreign key columns // in the specified ForeignKey table that reference the primary key or // columns representing a restraint of the parent table (could be the same // or a different table). Returns a FlightInfo object indicating where // the response can be retrieved with DoGet. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetCrossReference(ctx context.Context, pkTable, fkTable TableRef, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := pb.CommandGetCrossReference{ PkCatalog: pkTable.Catalog, @@ -227,32 +204,23 @@ func (c *Client) GetCrossReference(ctx context.Context, pkTable, fkTable TableRe return flightInfoForCommand(ctx, c, &cmd, opts...) } -<<<<<<< HEAD // GetTableTypes requests a list of the types of tables available on this // server. Returns a FlightInfo object indicating where the response can // be retrieved. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetTableTypes(ctx context.Context, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return flightInfoForCommand(ctx, c, &pb.CommandGetTableTypes{}, opts...) } -<<<<<<< HEAD // GetXdbcTypeInfo requests the information about all the data types supported // (dataType == nil) or a specific data type. Returns a FlightInfo object // indicating where the response can be retrieved. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetXdbcTypeInfo(ctx context.Context, dataType *int32, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return flightInfoForCommand(ctx, c, &pb.CommandGetXdbcTypeInfo{DataType: dataType}, opts...) } -<<<<<<< HEAD // GetSqlInfo returns a list of the requested SQL information corresponding // to the values in the info slice. Returns a FlightInfo object indicating // where the response can be retrieved. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.CallOption) (*flight.FlightInfo, error) { cmd := &pb.CommandGetSqlInfo{Info: make([]uint32, len(info))} @@ -262,6 +230,10 @@ func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.Ca return flightInfoForCommand(ctx, c, cmd, opts...) } +// Prepare creates a PreparedStatement object for the specified query. +// The resulting PreparedStatement object should be Closed when no longer +// needed. It will maintain a reference to this Client for use to execute +// and use the specified allocator for any allocations it needs to perform. // Prepare creates a PreparedStatement object for the specified query. // The resulting PreparedStatement object should be Closed when no longer // needed. It will maintain a reference to this Client for use to execute @@ -328,7 +300,6 @@ func (c *Client) Prepare(ctx context.Context, mem memory.Allocator, query string return } -<<<<<<< HEAD func (c *Client) getFlightInfo(ctx context.Context, desc *flight.FlightDescriptor, opts ...grpc.CallOption) (*flight.FlightInfo, error) { return c.Client.GetFlightInfo(ctx, desc, opts...) } @@ -343,14 +314,6 @@ func (c *Client) Close() error { return c.Client.Close() } // If the server returned the Dataset Schema or Parameter Binding schemas // at creation, they will also be accessible from this object. Close // should be called when no longer needed. -======= -func (c *Client) GetFlightInfo(ctx context.Context, desc *flight.FlightDescriptor, opts ...grpc.CallOption) (*flight.FlightInfo, error) { - return c.Client.GetFlightInfo(ctx, desc, opts...) -} - -func (c *Client) Close() error { return c.Client.Close() } - ->>>>>>> 7c0a7b52f (initial flightsql client) type PreparedStatement struct { client *Client opts []grpc.CallOption @@ -361,14 +324,11 @@ type PreparedStatement struct { closed bool } -<<<<<<< HEAD // Execute executes the prepared statement on the server and returns a FlightInfo // indicating where to retrieve the response. If SetParameters has been called // then the parameter bindings will be sent before execution. // // Will error if already closed. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) Execute(ctx context.Context) (*flight.FlightInfo, error) { if p.closed { return nil, errors.New("arrow/flightsql: prepared statement already closed") @@ -464,7 +424,6 @@ func (p *PreparedStatement) ExecuteUpdate(ctx context.Context) (nrecords int64, return updateResult.GetRecordCount(), nil } -<<<<<<< HEAD // DatasetSchema may be nil if the server did not return it when creating the // Prepared Statement. func (p *PreparedStatement) DatasetSchema() *arrow.Schema { return p.datasetSchema } @@ -480,11 +439,6 @@ func (p *PreparedStatement) ParameterSchema() *arrow.Schema { return p.paramSche // out from under the statement. Release will be called on a previous // binding record if it existed, and will be called upon calling Close // on the PreparedStatement. -======= -func (p *PreparedStatement) DatasetSchema() *arrow.Schema { return p.datasetSchema } -func (p *PreparedStatement) ParameterSchema() *arrow.Schema { return p.paramSchema } - ->>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) SetParameters(binding arrow.Record) { if p.paramBinding != nil { p.paramBinding.Release() @@ -494,12 +448,9 @@ func (p *PreparedStatement) SetParameters(binding arrow.Record) { p.paramBinding.Retain() } -<<<<<<< HEAD // Close calls release on any parameter binding record and sends // a ClosePreparedStatement action to the server. After calling // Close, the PreparedStatement should not be used again. -======= ->>>>>>> 7c0a7b52f (initial flightsql client) func (p *PreparedStatement) Close(ctx context.Context) error { if p.closed { return errors.New("arrow/flightsql: already closed") diff --git a/go/arrow/flight/flightsql/types.go b/go/arrow/flight/flightsql/types.go index 96bcf8073ef15..5e033d00ee322 100644 --- a/go/arrow/flight/flightsql/types.go +++ b/go/arrow/flight/flightsql/types.go @@ -22,6 +22,7 @@ import ( "google.golang.org/protobuf/types/known/anypb" ) +// Constants for Action types const ( CreatePreparedStatementActionType = "CreatePreparedStatement" ClosePreparedStatementActionType = "ClosePreparedStatement" @@ -66,6 +67,9 @@ func impkToTableRef(cmd *pb.CommandGetImportedKeys) TableRef { } } +// CreateStatementQueryTicket is a helper that constructs a properly +// serialized TicketStatementQuery containing a given opaque binary handle +// for use with constructing a ticket to return from GetFlightInfoStatement. func CreateStatementQueryTicket(handle []byte) ([]byte, error) { query := &pb.TicketStatementQuery{StatementHandle: handle} var ticket anypb.Any @@ -75,11 +79,22 @@ func CreateStatementQueryTicket(handle []byte) ([]byte, error) { } type ( - GetDBSchemasOpts = pb.CommandGetDbSchemas - GetTablesOpts = pb.CommandGetTables - + // GetDBSchemasOpts contains the options to request Database Schemas: + // an optional Catalog and a Schema Name filter pattern. + GetDBSchemasOpts pb.CommandGetDbSchemas + // GetTablesOpts contains the options for retrieving a list of tables: + // optional Catalog, Schema filter pattern, Table name filter pattern, + // a filter of table types, and whether or not to include the schema + // in the response. + GetTablesOpts pb.CommandGetTables + + // SqlInfoResultMap is a mapping of SqlInfo ids to the desired response. + // This is part of a Server and used for registering responses to a + // SqlInfo request. SqlInfoResultMap map[uint32]interface{} + // TableRef is a helpful struct for referencing a specific Table + // by its catalog, schema, and table name. TableRef struct { // Catalog specifies the catalog this table belongs to. // An empty string refers to tables without a catalog. @@ -94,6 +109,8 @@ type ( Table string } + // CrossTableRef contains a reference to a Primary Key table + // and a Foreign Key table. CrossTableRef struct { PKRef TableRef FKRef TableRef From 55a6e4b8f0f121d2182dfacb818012d5c44bd350 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 11 Aug 2022 14:10:16 -0400 Subject: [PATCH 09/23] so far so good --- go/arrow/flight/flightsql/column_metadata.go | 5 + .../flightsql/example/sql_batch_reader.go | 236 +++++++++++++ .../flight/flightsql/example/sql_statement.go | 36 ++ .../flight/flightsql/example/sqlite_server.go | 327 ++++++++++++++++++ .../sqlite_tables_schema_batch_reader.go | 200 +++++++++++ .../flight/flightsql/example/sqlite_test.go | 120 +++++++ go/arrow/flight/record_batch_reader.go | 13 +- go/go.mod | 10 +- go/go.sum | 56 +++ 9 files changed, 1001 insertions(+), 2 deletions(-) create mode 100644 go/arrow/flight/flightsql/example/sql_batch_reader.go create mode 100644 go/arrow/flight/flightsql/example/sql_statement.go create mode 100644 go/arrow/flight/flightsql/example/sqlite_server.go create mode 100644 go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go create mode 100644 go/arrow/flight/flightsql/example/sqlite_test.go diff --git a/go/arrow/flight/flightsql/column_metadata.go b/go/arrow/flight/flightsql/column_metadata.go index 59c133757d5db..6ad8030d9950f 100644 --- a/go/arrow/flight/flightsql/column_metadata.go +++ b/go/arrow/flight/flightsql/column_metadata.go @@ -142,6 +142,11 @@ func NewColumnMetadataBuilder() *ColumnMetadataBuilder { return &ColumnMetadataBuilder{make([]string, 0), make([]string, 0)} } +func (c *ColumnMetadataBuilder) Clear() { + c.keys = c.keys[:0] + c.vals = c.vals[:0] +} + func (c *ColumnMetadataBuilder) Build() ColumnMetadata { md := c.Metadata() return ColumnMetadata{&md} diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go new file mode 100644 index 0000000000000..e142e55e06878 --- /dev/null +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -0,0 +1,236 @@ +// 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. + +package example + +import ( + "database/sql" + "strings" + "sync/atomic" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/internal/debug" + "github.com/apache/arrow/go/v10/arrow/memory" +) + +func getArrowType(dbtype string) arrow.DataType { + dbtype = strings.ToLower(dbtype) + + if strings.HasPrefix(dbtype, "varchar") { + return arrow.BinaryTypes.String + } + + switch dbtype { + case "int", "integer": + return arrow.PrimitiveTypes.Int64 + case "real": + return arrow.PrimitiveTypes.Float64 + case "blob": + return arrow.BinaryTypes.Binary + case "text", "date", "char": + return arrow.BinaryTypes.String + default: + panic("invalid sqlite type: " + dbtype) + } +} + +// func getColumnMetadata(bldr *flightsql.ColumnMetadataBuilder, c *sql.ColumnType) arrow.Metadata { +// defer bldr.Clear() +// bldr.IsAutoIncrement(false).IsReadOnly(false) +// prec, scale, ok := c.DecimalSize() +// if ok { +// bldr.Precision(int32(prec)).Scale(int32(scale)) +// } +// return bldr.Metadata() +// } + +const maxBatchSize = 1024 + +type SqlBatchReader struct { + refCount int64 + + schema *arrow.Schema + rows *sql.Rows + record arrow.Record + bldr *array.RecordBuilder + err error + + rowdest []interface{} +} + +func NewSqlBatchReaderWithSchema(mem memory.Allocator, schema *arrow.Schema, rows *sql.Rows) (*SqlBatchReader, error) { + rowdest := make([]interface{}, len(schema.Fields())) + for i, f := range schema.Fields() { + switch f.Type.ID() { + case arrow.INT64: + if f.Nullable { + rowdest[i] = &sql.NullInt64{} + } else { + rowdest[i] = new(int64) + } + case arrow.FLOAT64: + if f.Nullable { + rowdest[i] = &sql.NullFloat64{} + } else { + rowdest[i] = new(float64) + } + case arrow.BINARY: + var b []byte + rowdest[i] = &b + case arrow.STRING: + if f.Nullable { + rowdest[i] = &sql.NullString{} + } else { + rowdest[i] = new(string) + } + } + } + + return &SqlBatchReader{ + refCount: 1, + bldr: array.NewRecordBuilder(mem, schema), + schema: schema, + rowdest: rowdest, + rows: rows}, nil +} + +func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, error) { + cols, err := rows.ColumnTypes() + if err != nil { + rows.Close() + return nil, err + } + + rowdest := make([]interface{}, len(cols)) + fields := make([]arrow.Field, len(cols)) + for i, c := range cols { + fields[i].Name = c.Name() + fields[i].Nullable, _ = c.Nullable() + fields[i].Type = getArrowType(c.DatabaseTypeName()) + switch fields[i].Type.ID() { + case arrow.INT64: + if fields[i].Nullable { + rowdest[i] = &sql.NullInt64{} + } else { + rowdest[i] = new(int64) + } + case arrow.FLOAT64: + if fields[i].Nullable { + rowdest[i] = &sql.NullFloat64{} + } else { + rowdest[i] = new(float64) + } + case arrow.BINARY: + var b []byte + rowdest[i] = &b + case arrow.STRING: + if fields[i].Nullable { + rowdest[i] = &sql.NullString{} + } else { + rowdest[i] = new(string) + } + } + } + + schema := arrow.NewSchema(fields, nil) + return &SqlBatchReader{ + refCount: 1, + bldr: array.NewRecordBuilder(mem, schema), + schema: schema, + rowdest: rowdest, + rows: rows}, nil +} + +func (r *SqlBatchReader) Retain() { + atomic.AddInt64(&r.refCount, 1) +} + +func (r *SqlBatchReader) Release() { + debug.Assert(atomic.LoadInt64(&r.refCount) > 0, "too many releases") + + if atomic.AddInt64(&r.refCount, -1) == 0 { + r.rows.Close() + r.rows, r.schema, r.rowdest = nil, nil, nil + r.bldr.Release() + r.bldr = nil + if r.record != nil { + r.record.Release() + r.record = nil + } + } +} +func (r *SqlBatchReader) Schema() *arrow.Schema { return r.schema } + +func (r *SqlBatchReader) Record() arrow.Record { return r.record } + +func (r *SqlBatchReader) Err() error { return r.err } + +func (r *SqlBatchReader) Next() bool { + if r.record != nil { + r.record.Release() + r.record = nil + } + + rows := 0 + for rows < maxBatchSize && r.rows.Next() { + if err := r.rows.Scan(r.rowdest...); err != nil { + r.err = err + return false + } + + for i, v := range r.rowdest { + fb := r.bldr.Field(i) + switch v := v.(type) { + case *int64: + fb.(*array.Int64Builder).Append(*v) + case *sql.NullInt64: + if !v.Valid { + fb.AppendNull() + } else { + fb.(*array.Int64Builder).Append(v.Int64) + } + case *float64: + fb.(*array.Float64Builder).Append(*v) + case *sql.NullFloat64: + if !v.Valid { + fb.AppendNull() + } else { + fb.(*array.Float64Builder).Append(v.Float64) + } + case *[]byte: + if v == nil { + fb.AppendNull() + } else { + fb.(*array.BinaryBuilder).Append(*v) + } + case *string: + fb.(*array.StringBuilder).Append(*v) + case *sql.NullString: + if !v.Valid { + fb.AppendNull() + } else { + fb.(*array.StringBuilder).Append(v.String) + } + } + } + + rows++ + } + + r.record = r.bldr.NewRecord() + return rows > 0 +} diff --git a/go/arrow/flight/flightsql/example/sql_statement.go b/go/arrow/flight/flightsql/example/sql_statement.go new file mode 100644 index 0000000000000..6ae75a55d9122 --- /dev/null +++ b/go/arrow/flight/flightsql/example/sql_statement.go @@ -0,0 +1,36 @@ +// 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. + +package example + +import ( + "database/sql" +) + +// func GetColumnMetadata(coltype int, table string) flightsql.ColumnMetadata {} + +type Statement struct { + stmt *sql.Stmt + db *sql.DB +} + +func NewStatement(db *sql.DB, query string) (*Statement, error) { + stmt, err := db.Prepare(query) + if err != nil { + return nil, err + } + return &Statement{stmt, db}, nil +} diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go new file mode 100644 index 0000000000000..47dc80de5c9c1 --- /dev/null +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -0,0 +1,327 @@ +// 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. + +// Package example contains a FlightSQL Server implementation using +// sqlite as the backing engine. +// +// In order to ensure portability we'll use modernc.org/sqlite instead +// of github.com/mattn/go-sqlite3 because modernc is a translation of the +// SQLite source into Go, such that it doesn't require CGO to run and +// doesn't need to link against the actual libsqlite3 libraries. This way +// we don't require CGO or libsqlite3 to run this example or the tests. +// +// That said, since both implement in terms of Go's standard database/sql +// package, it's easy to swap them out if desired as the modernc.org/sqlite +// package is slower than go-sqlite3. +package example + +import ( + "context" + "database/sql" + "fmt" + "math/rand" + "strings" + "sync" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" + "github.com/apache/arrow/go/v10/arrow/memory" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + _ "modernc.org/sqlite" +) + +func genRandomString() []byte { + const length = 16 + max := int('z') + min := int('0') + + out := make([]byte, length) + for i := range out { + out[i] = byte(rand.Intn(max-min+1) + min) + } + return out +} + +func prepareQueryForGetTables(cmd flightsql.GetTables) string { + var b strings.Builder + b.WriteString(`SELECT null AS catalog_name, null AS schema_name, + name AS table_name, type AS table_type FROM sqlite_master WHERE 1=1`) + + if cmd.GetCatalog() != nil { + b.WriteString(" and catalog_name = '") + b.WriteString(*cmd.GetCatalog()) + b.WriteByte('\'') + } + + if cmd.GetDBSchemaFilterPattern() != nil { + b.WriteString(" and schema_name LIKE '") + b.WriteString(*cmd.GetDBSchemaFilterPattern()) + b.WriteByte('\'') + } + + if cmd.GetTableNameFilterPattern() != nil { + b.WriteString(" and table_name LIKE '") + b.WriteString(*cmd.GetTableNameFilterPattern()) + b.WriteByte('\'') + } + + if len(cmd.GetTableTypes()) > 0 { + b.WriteString(" and table_type IN (") + for i, t := range cmd.GetTableTypes() { + if i != 0 { + b.WriteByte(',') + } + fmt.Fprintf(&b, "'%s'", t) + } + b.WriteByte(')') + } + + b.WriteString(" order by table_name") + return b.String() +} + +type SQLiteFlightSQLServer struct { + flightsql.BaseServer + db *sql.DB + mem memory.Allocator + + prepared sync.Map +} + +func NewSQLiteFlightSQLServer() (*SQLiteFlightSQLServer, error) { + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + return nil, err + } + + _, err = db.Exec(` + CREATE TABLE foreignTable ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + foreignName varchar(100), + value int); + + CREATE TABLE intTable ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + keyName varchar(100), + value int, + foreignId int references foreignTable(id)); + + INSERT INTO foreignTable (foreignName, value) VALUES ('keyOne', 1); + INSERT INTO foreignTable (foreignName, value) VALUES ('keyTwo', 0); + INSERT INTO foreignTable (foreignName, value) VALUES ('keyThree', -1); + INSERT INTO intTable (keyName, value, foreignId) VALUES ('one', 1, 1); + INSERT INTO intTable (keyName, value, foreignId) VALUES ('zero', 0, 1); + INSERT INTO intTable (keyName, value, foreignId) VALUES ('negative one', -1, 1); + INSERT INTO intTable (keyName, value, foreignId) VALUES (NULL, NULL, NULL); + `) + + if err != nil { + return nil, err + } + return &SQLiteFlightSQLServer{db: db}, nil +} + +func (s *SQLiteFlightSQLServer) flightInfoForCommand(desc *flight.FlightDescriptor, schema *arrow.Schema) *flight.FlightInfo { + return &flight.FlightInfo{ + Endpoint: []*flight.FlightEndpoint{{Ticket: &flight.Ticket{Ticket: desc.Cmd}}}, + FlightDescriptor: desc, + Schema: flight.SerializeSchema(schema, s.mem), + TotalRecords: -1, + TotalBytes: -1, + } +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoStatement(ctx context.Context, cmd flightsql.StatementQuery, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + query := cmd.GetQuery() + tkt, err := flightsql.CreateStatementQueryTicket([]byte(query)) + if err != nil { + return nil, err + } + + return &flight.FlightInfo{ + Endpoint: []*flight.FlightEndpoint{{Ticket: &flight.Ticket{Ticket: tkt}}}, + FlightDescriptor: desc, + TotalRecords: -1, + TotalBytes: -1, + }, nil +} + +func (s *SQLiteFlightSQLServer) DoGetStatement(ctx context.Context, cmd flightsql.StatementQueryTicket) (*arrow.Schema, <-chan flight.StreamChunk, error) { + rows, err := s.db.QueryContext(ctx, string(cmd.GetStatementHandle())) + if err != nil { + return nil, nil, err + } + + reader, err := NewSqlBatchReader(s.mem, rows) + if err != nil { + return nil, nil, err + } + + ch := make(chan flight.StreamChunk) + go flight.StreamChunksFromReader(reader, ch) + return reader.schema, ch, nil +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoCatalogs(_ context.Context, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.Catalogs), nil +} + +func (s *SQLiteFlightSQLServer) DoGetCatalogs(context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { + // sqlite doesn't support catalogs, this returns an empty record batch + schema := schema_ref.Catalogs + batchBldr := array.NewRecordBuilder(s.mem, schema) + defer batchBldr.Release() + + ch := make(chan flight.StreamChunk, 1) + ch <- flight.StreamChunk{Data: batchBldr.NewRecord()} + close(ch) + + return schema, ch, nil +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoSchemas(_ context.Context, cmd flightsql.GetDBSchemas, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.DBSchemas), nil +} + +func (s *SQLiteFlightSQLServer) DoGetDBSchemas(context.Context, flightsql.GetDBSchemas) (*arrow.Schema, <-chan flight.StreamChunk, error) { + // sqlite doesn't support schemas, this returns an empty record batch + schema := schema_ref.DBSchemas + batchBldr := array.NewRecordBuilder(s.mem, schema) + defer batchBldr.Release() + + ch := make(chan flight.StreamChunk, 1) + ch <- flight.StreamChunk{Data: batchBldr.NewRecord()} + close(ch) + + return schema, ch, nil +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoTables(_ context.Context, cmd flightsql.GetTables, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + schema := schema_ref.Tables + if cmd.GetIncludeSchema() { + schema = schema_ref.TablesWithIncludedSchema + } + return s.flightInfoForCommand(desc, schema), nil +} + +func (s *SQLiteFlightSQLServer) DoGetTables(ctx context.Context, cmd flightsql.GetTables) (*arrow.Schema, <-chan flight.StreamChunk, error) { + query := prepareQueryForGetTables(cmd) + + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return nil, nil, err + } + + var rdr array.RecordReader + + rdr, err = NewSqlBatchReaderWithSchema(s.mem, schema_ref.Tables, rows) + if err != nil { + return nil, nil, err + } + + ch := make(chan flight.StreamChunk, 2) + if cmd.GetIncludeSchema() { + rdr, err = NewSqliteTablesSchemaBatchReader(ctx, s.mem, rdr, s.db, query) + if err != nil { + return nil, nil, err + } + } + + go flight.StreamChunksFromReader(rdr, ch) + return rdr.Schema(), ch, nil +} + +func (s *SQLiteFlightSQLServer) DoPutCommandStatementUpdate(ctx context.Context, cmd flightsql.StatementUpdate) (int64, error) { + res, err := s.db.ExecContext(ctx, cmd.GetQuery()) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + +func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req flightsql.ActionCreatePreparedStatementRequest) (result flightsql.ActionCreatePreparedStatementResult, err error) { + stmt, err := s.db.PrepareContext(ctx, req.GetQuery()) + if err != nil { + return result, err + } + + handle := genRandomString() + s.prepared.Store(handle, stmt) + + result.Handle = handle + // no way to get the dataset or parameter schemas from sql.DB + return +} + +func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, request flightsql.ActionClosePreparedStatementRequest) error { + handle := request.GetPreparedStatementHandle() + if val, loaded := s.prepared.LoadAndDelete(handle); loaded { + stmt := val.(sql.Stmt) + return stmt.Close() + } + + return status.Error(codes.InvalidArgument, "prepared statement not found") +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoPreparedStatement(_ context.Context, cmd flightsql.PreparedStatementQuery, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + _, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) + if !ok { + return nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } + + return &flight.FlightInfo{ + Endpoint: []*flight.FlightEndpoint{{Ticket: &flight.Ticket{Ticket: desc.Cmd}}}, + FlightDescriptor: desc, + TotalRecords: -1, + TotalBytes: -1, + }, nil +} + +func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd flightsql.PreparedStatementQuery) (*arrow.Schema, <-chan flight.StreamChunk, error) { + val, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) + if !ok { + return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") + } + + stmt := val.(sql.Stmt) + rows, err := stmt.QueryContext(ctx) + if err != nil { + return nil, nil, err + } + + rdr, err := NewSqlBatchReader(s.mem, rows) + if err != nil { + return nil, nil, err + } + + ch := make(chan flight.StreamChunk) + go flight.StreamChunksFromReader(rdr, ch) + return rdr.Schema(), ch, nil +} + +// func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd flightsql.PreparedStatementUpdate, rdr flight.MessageReader) (int64, error) { +// val, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) +// if !ok { +// return 0, status.Error(codes.InvalidArgument, "prepared statement not found") +// } + +// stmt := val.(sql.Stmt) + +// } diff --git a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go new file mode 100644 index 0000000000000..90d9bde94e73c --- /dev/null +++ b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go @@ -0,0 +1,200 @@ +// 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. + +package example + +import ( + "context" + "database/sql" + "strings" + "sync/atomic" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + "github.com/apache/arrow/go/v10/arrow/internal/debug" + "github.com/apache/arrow/go/v10/arrow/memory" + sqlite3 "modernc.org/sqlite/lib" +) + +type SqliteTablesSchemaBatchReader struct { + refCount int64 + + mem memory.Allocator + ctx context.Context + rdr array.RecordReader + stmt *sql.Stmt + schemaBldr *array.BinaryBuilder + record arrow.Record + err error +} + +func NewSqliteTablesSchemaBatchReader(ctx context.Context, mem memory.Allocator, rdr array.RecordReader, db *sql.DB, mainQuery string) (*SqliteTablesSchemaBatchReader, error) { + schemaQuery := `SELECT table_name, name, type, [notnull] + FROM pragma_table_info(table_name) + JOIN (` + mainQuery + `) WHERE table_name = ?` + + stmt, err := db.PrepareContext(ctx, schemaQuery) + if err != nil { + rdr.Release() + return nil, err + } + + return &SqliteTablesSchemaBatchReader{ + refCount: 1, + ctx: ctx, + rdr: rdr, + stmt: stmt, + mem: mem, + schemaBldr: array.NewBinaryBuilder(mem, arrow.BinaryTypes.Binary), + }, nil +} + +func (s *SqliteTablesSchemaBatchReader) Err() error { return s.err } + +func (s *SqliteTablesSchemaBatchReader) Retain() { atomic.AddInt64(&s.refCount, 1) } + +func (s *SqliteTablesSchemaBatchReader) Release() { + debug.Assert(atomic.LoadInt64(&s.refCount) > 0, "too many releases") + + if atomic.AddInt64(&s.refCount, -1) == 0 { + s.rdr.Release() + s.stmt.Close() + s.schemaBldr.Release() + if s.record != nil { + s.record.Release() + s.record = nil + } + } +} + +func (s *SqliteTablesSchemaBatchReader) Schema() *arrow.Schema { + fields := append(s.rdr.Schema().Fields(), + arrow.Field{Name: "table_schema", Type: arrow.BinaryTypes.Binary}) + return arrow.NewSchema(fields, nil) +} + +func (s *SqliteTablesSchemaBatchReader) Record() arrow.Record { return s.record } + +func getSqlTypeFromTypeName(sqltype string) int { + if sqltype == "" { + return sqlite3.SQLITE_NULL + } + + sqltype = strings.ToLower(sqltype) + + if strings.HasPrefix(sqltype, "varchar") || strings.HasPrefix(sqltype, "char") { + return sqlite3.SQLITE_TEXT + } + + switch sqltype { + case "int", "integer": + return sqlite3.SQLITE_INTEGER + case "real": + return sqlite3.SQLITE_FLOAT + case "blob": + return sqlite3.SQLITE_BLOB + case "text", "date": + return sqlite3.SQLITE_TEXT + default: + return sqlite3.SQLITE_NULL + } +} + +func getPrecisionFromCol(sqltype int) int { + switch sqltype { + case sqlite3.SQLITE_INTEGER: + return 10 + case sqlite3.SQLITE_FLOAT: + return 15 + } + return 0 +} + +func getColumnMetadata(bldr *flightsql.ColumnMetadataBuilder, sqltype int, table string) arrow.Metadata { + defer bldr.Clear() + + bldr.Scale(15).IsReadOnly(false) + if table != "" { + bldr.TableName(table) + } + switch sqltype { + case sqlite3.SQLITE_TEXT, sqlite3.SQLITE_BLOB: + default: + bldr.Precision(int32(getPrecisionFromCol(sqltype))) + } + + return bldr.Metadata() +} + +func (s *SqliteTablesSchemaBatchReader) Next() bool { + if s.record != nil { + s.record.Release() + s.record = nil + } + + if !s.rdr.Next() { + return false + } + + rec := s.rdr.Record() + tableNameArr := rec.Column(rec.Schema().FieldIndices("table_name")[0]).(*array.String) + + bldr := flightsql.NewColumnMetadataBuilder() + columnFields := make([]arrow.Field, 0) + for i := 0; i < tableNameArr.Len(); i++ { + table := tableNameArr.Value(i) + rows, err := s.stmt.QueryContext(s.ctx, table) + if err != nil { + s.err = err + return false + } + + var tableName, name, typ string + var nn int + for rows.Next() { + if err := rows.Scan(&tableName, &name, &typ, &nn); err != nil { + rows.Close() + s.err = err + return false + } + + columnFields = append(columnFields, arrow.Field{ + Name: name, + Type: getArrowType(typ), + Nullable: nn == 1, + Metadata: getColumnMetadata(bldr, getSqlTypeFromTypeName(typ), tableName), + }) + } + + rows.Close() + if rows.Err() != nil { + s.err = rows.Err() + return false + } + val := flight.SerializeSchema(arrow.NewSchema(columnFields, nil), s.mem) + s.schemaBldr.Append(val) + + columnFields = columnFields[:0] + } + + schemaCol := s.schemaBldr.NewArray() + defer schemaCol.Release() + + s.record = array.NewRecord(s.Schema(), append(rec.Columns(), schemaCol), rec.NumRows()) + return true +} diff --git a/go/arrow/flight/flightsql/example/sqlite_test.go b/go/arrow/flight/flightsql/example/sqlite_test.go new file mode 100644 index 0000000000000..52c107ff83bb3 --- /dev/null +++ b/go/arrow/flight/flightsql/example/sqlite_test.go @@ -0,0 +1,120 @@ +// 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. + +package example + +import ( + "context" + "strings" + "testing" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" + "github.com/apache/arrow/go/v10/arrow/memory" + "github.com/apache/arrow/go/v10/arrow/scalar" + "github.com/stretchr/testify/assert" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +var dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} + +func TestServer(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + s := flight.NewServerWithMiddleware(nil) + srv, err := NewSQLiteFlightSQLServer() + assert.NoError(t, err) + assert.NotNil(t, srv) + srv.mem = mem + s.RegisterFlightService(flightsql.NewFlightServer(srv)) + s.Init("localhost:0") + + go s.Serve() + defer s.Shutdown() + + cl, err := flightsql.NewClient(s.Addr().String(), nil, nil, dialOpts...) + assert.NoError(t, err) + assert.NotNil(t, cl) + + info, err := cl.Execute(context.Background(), "SELECT * FROM intTable") + assert.NoError(t, err) + + rdr, err := cl.DoGet(context.Background(), info.Endpoint[0].Ticket) + assert.NoError(t, err) + defer rdr.Release() + + assert.True(t, rdr.Next()) + rec := rdr.Record() + assert.NotNil(t, rec) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + }, nil) + + assert.Truef(t, expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) + + idarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 2, 3, 4]`)) + defer idarr.Release() + keyarr, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["one", "zero", "negative one", null]`)) + defer keyarr.Release() + valarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 0, -1, null]`)) + defer valarr.Release() + foreignarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 1, 1, null]`)) + defer foreignarr.Release() + + expectedRec := array.NewRecord(expectedSchema, []arrow.Array{idarr, keyarr, valarr, foreignarr}, 4) + defer expectedRec.Release() + + assert.Truef(t, array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) + + info, err = cl.GetTables(context.Background(), &flightsql.GetTablesOpts{}) + assert.NoError(t, err) + assert.NotNil(t, info) + + rdr, err = cl.DoGet(context.Background(), info.Endpoint[0].Ticket) + assert.NoError(t, err) + defer rdr.Release() + + catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, mem) + defer catalogName.Release() + schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, mem) + defer schemaName.Release() + + tableName, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["foreignTable", "intTable", "sqlite_sequence"]`)) + defer tableName.Release() + + tableType, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["table", "table", "table"]`)) + defer tableType.Release() + + expectedRec = array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) + defer expectedRec.Release() + + assert.True(t, rdr.Next()) + rec = rdr.Record() + assert.NotNil(t, rec) + rec.Retain() + assert.False(t, rdr.Next()) + + assert.Truef(t, array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} diff --git a/go/arrow/flight/record_batch_reader.go b/go/arrow/flight/record_batch_reader.go index 75e09f2008f6b..035ba9c4bbeff 100644 --- a/go/arrow/flight/record_batch_reader.go +++ b/go/arrow/flight/record_batch_reader.go @@ -153,12 +153,13 @@ func NewRecordReader(r DataStreamReader, opts ...ipc.Option) (*Reader, error) { return nil, err } - rdr := &Reader{dmr: &dataMessageReader{rdr: r}} + rdr := &Reader{dmr: &dataMessageReader{rdr: r, refCount: 1}} rdr.dmr.descr = data.FlightDescriptor if len(data.DataHeader) > 0 { rdr.dmr.peeked = data } + rdr.dmr.Retain() if rdr.Reader, err = ipc.NewReaderFromMessageReader(rdr.dmr, opts...); err != nil { return nil, fmt.Errorf("arrow/flight: could not create flight reader: %w", err) } @@ -201,6 +202,10 @@ type MessageReader interface { LatestAppMetadata() []byte } +type haserr interface { + Err() error +} + // StreamChunksFromReader is a convenience function to populate a channel // from a record reader. It is intended to be run using a separate goroutine // by calling `go flight.StreamChunksFromReader(rdr, ch)`. @@ -222,4 +227,10 @@ func StreamChunksFromReader(rdr array.RecordReader, ch chan<- StreamChunk) { rec.Retain() ch <- StreamChunk{Data: rec} } + + if e, ok := rdr.(haserr); ok { + if e.Err() != nil { + ch <- StreamChunk{Err: e.Err()} + } + } } diff --git a/go/go.mod b/go/go.mod index 18a1c9abfa773..6a41ad0f06a17 100644 --- a/go/go.mod +++ b/go/go.mod @@ -30,6 +30,7 @@ require ( github.com/klauspost/asmfmt v1.3.2 github.com/klauspost/compress v1.15.9 github.com/kr/pretty v0.1.0 // indirect + github.com/mattn/go-isatty v0.0.14 // indirect github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 github.com/pierrec/lz4/v4 v4.1.15 @@ -37,11 +38,18 @@ require ( github.com/zeebo/xxh3 v1.0.2 golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4 - golang.org/x/sys v0.0.0-20220804214406-8e32c043e418 + golang.org/x/sys v0.0.0-20220808155132-1c4a2a72c664 golang.org/x/tools v0.1.12 golang.org/x/xerrors v0.0.0-20220609144429-65e65417b02f gonum.org/v1/gonum v0.11.0 google.golang.org/grpc v1.48.0 google.golang.org/protobuf v1.28.1 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect + lukechampine.com/uint128 v1.2.0 // indirect + modernc.org/cc/v3 v3.36.1 // indirect + modernc.org/ccgo/v3 v3.16.8 // indirect + modernc.org/libc v1.16.19 // indirect + modernc.org/opt v0.1.3 // indirect + modernc.org/sqlite v1.18.0 // indirect + modernc.org/strutil v1.1.2 // indirect ) diff --git a/go/go.sum b/go/go.sum index 7653c2cc1b917..c2973d283fa67 100644 --- a/go/go.sum +++ b/go/go.sum @@ -32,6 +32,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815 h1:bWDMxwH3px2JBh6AyO7hdCn/PkvCZXii8TGj7sbtEbQ= github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE= +github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= @@ -80,15 +81,20 @@ github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMyw github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE= github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg= github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw= github.com/jung-kurt/gofpdf v1.0.0/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes= github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4= github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE= @@ -101,6 +107,10 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= +github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= +github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= +github.com/mattn/go-sqlite3 v1.14.12/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI= @@ -115,6 +125,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 h1:OdAsTTz6OkFY5QxjkYwrChwuRruF69c169dPK26NUlk= +github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ= github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w= github.com/ruudk/golang-pdf417 v0.0.0-20201230142125-a7e3863a1245/go.mod h1:pQAZKsJ8yyVxGRWYNEm9oFB8ieLgKFnamEyDmSA0BRk= @@ -198,6 +210,7 @@ golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5h golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -205,11 +218,15 @@ golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210304124612-50617c2ba197/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20211019181941-9d821ace8654/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220804214406-8e32c043e418 h1:9vYwv7OjYaky/tlAeD7C4oC9EsPTlaFl1H2jS++V+ME= golang.org/x/sys v0.0.0-20220804214406-8e32c043e418/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220808155132-1c4a2a72c664 h1:v1W7bwXHsnLLloWYTVEdvGvA7BHMeBYsPcF0GLDxIRs= +golang.org/x/sys v0.0.0-20220808155132-1c4a2a72c664/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= @@ -227,6 +244,7 @@ golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3 golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= golang.org/x/tools v0.0.0-20190927191325-030b2cf1153e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201124115921-2c860bdd6e78/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= golang.org/x/tools v0.1.9/go.mod h1:nABZi5QlRsZVlzPpHl034qft6wpY4eDcsTt5AaioBiU= golang.org/x/tools v0.1.10/go.mod h1:Uh6Zz+xoGYZom868N8YTex3t7RhtHDBrE8Gzo9bV56E= @@ -287,4 +305,42 @@ gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.1.3/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las= +lukechampine.com/uint128 v1.1.1/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl2j1gk= +lukechampine.com/uint128 v1.2.0 h1:mBi/5l91vocEN8otkC5bDLhi2KdCticRiwbdB0O+rjI= +lukechampine.com/uint128 v1.2.0/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl2j1gk= +modernc.org/cc/v3 v3.36.0/go.mod h1:NFUHyPn4ekoC/JHeZFfZurN6ixxawE1BnVonP/oahEI= +modernc.org/cc/v3 v3.36.1 h1:CICrjwr/1M4+6OQ4HJZ/AHxjcwe67r5vPUF518MkO8A= +modernc.org/cc/v3 v3.36.1/go.mod h1:NFUHyPn4ekoC/JHeZFfZurN6ixxawE1BnVonP/oahEI= +modernc.org/ccgo/v3 v3.0.0-20220428102840-41399a37e894/go.mod h1:eI31LL8EwEBKPpNpA4bU1/i+sKOwOrQy8D87zWUcRZc= +modernc.org/ccgo/v3 v3.0.0-20220430103911-bc99d88307be/go.mod h1:bwdAnOoaIt8Ax9YdWGjxWsdkPcZyRPHqrOvJxaKAKGw= +modernc.org/ccgo/v3 v3.16.4/go.mod h1:tGtX0gE9Jn7hdZFeU88slbTh1UtCYKusWOoCJuvkWsQ= +modernc.org/ccgo/v3 v3.16.6/go.mod h1:tGtX0gE9Jn7hdZFeU88slbTh1UtCYKusWOoCJuvkWsQ= +modernc.org/ccgo/v3 v3.16.8 h1:G0QNlTqI5uVgczBWfGKs7B++EPwCfXPWGD2MdeKloDs= +modernc.org/ccgo/v3 v3.16.8/go.mod h1:zNjwkizS+fIFDrDjIAgBSCLkWbJuHF+ar3QRn+Z9aws= +modernc.org/ccorpus v1.11.6/go.mod h1:2gEUTrWqdpH2pXsmTM1ZkjeSrUWDpjMu2T6m29L/ErQ= +modernc.org/httpfs v1.0.6/go.mod h1:7dosgurJGp0sPaRanU53W4xZYKh14wfzX420oZADeHM= +modernc.org/libc v0.0.0-20220428101251-2d5f3daf273b/go.mod h1:p7Mg4+koNjc8jkqwcoFBJx7tXkpj00G77X7A72jXPXA= +modernc.org/libc v1.16.0/go.mod h1:N4LD6DBE9cf+Dzf9buBlzVJndKr/iJHG97vGLHYnb5A= +modernc.org/libc v1.16.1/go.mod h1:JjJE0eu4yeK7tab2n4S1w8tlWd9MxXLRzheaRnAKymU= +modernc.org/libc v1.16.7/go.mod h1:hYIV5VZczAmGZAnG15Vdngn5HSF5cSkbvfz2B7GRuVU= +modernc.org/libc v1.16.17/go.mod h1:hYIV5VZczAmGZAnG15Vdngn5HSF5cSkbvfz2B7GRuVU= +modernc.org/libc v1.16.19 h1:S8flPn5ZeXx6iw/8yNa986hwTQDrY8RXU7tObZuAozo= +modernc.org/libc v1.16.19/go.mod h1:p7Mg4+koNjc8jkqwcoFBJx7tXkpj00G77X7A72jXPXA= +modernc.org/mathutil v1.2.2/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= +modernc.org/mathutil v1.4.1 h1:ij3fYGe8zBF4Vu+g0oT7mB06r8sqGWKuJu1yXeR4by8= +modernc.org/mathutil v1.4.1/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= +modernc.org/memory v1.1.1 h1:bDOL0DIDLQv7bWhP3gMvIrnoFw+Eo6F7a2QK9HPDiFU= +modernc.org/memory v1.1.1/go.mod h1:/0wo5ibyrQiaoUoH7f9D8dnglAmILJ5/cxZlRECf+Nw= +modernc.org/opt v0.1.1/go.mod h1:WdSiB5evDcignE70guQKxYUl14mgWtbClRi5wmkkTX0= +modernc.org/opt v0.1.3 h1:3XOZf2yznlhC+ibLltsDGzABUGVx8J6pnFMS3E4dcq4= +modernc.org/opt v0.1.3/go.mod h1:WdSiB5evDcignE70guQKxYUl14mgWtbClRi5wmkkTX0= +modernc.org/sqlite v1.18.0 h1:ef66qJSgKeyLyrF4kQ2RHw/Ue3V89fyFNbGL073aDjI= +modernc.org/sqlite v1.18.0/go.mod h1:B9fRWZacNxJBHoCJZQr1R54zhVn3fjfl0aszflrTSxY= +modernc.org/strutil v1.1.1/go.mod h1:DE+MQQ/hjKBZS2zNInV5hhcipt5rLPWkmpbGeW5mmdw= +modernc.org/strutil v1.1.2 h1:iFBDH6j1Z0bN/Q9udJnnFoFpENA4252qe/7/5woE5MI= +modernc.org/strutil v1.1.2/go.mod h1:OYajnUAcI/MX+XD/Wx7v1bbdvcQSvxgtb0gC+u3d3eg= +modernc.org/tcl v1.13.1/go.mod h1:XOLfOwzhkljL4itZkK6T72ckMgvj0BDsnKNdZVUOecw= +modernc.org/token v1.0.0 h1:a0jaWiNMDhDUtqOj09wvjWWAqd3q7WpBulmL9H2egsk= +modernc.org/token v1.0.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= +modernc.org/z v1.5.1/go.mod h1:eWFB510QWW5Th9YGZT81s+LwvaAs3Q2yr4sP0rmLkv8= rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= From 0a7b3509a475c2111f30abfabad4ff9ae94747ea Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 12 Aug 2022 11:06:06 -0400 Subject: [PATCH 10/23] first batch of tests using sqlite example --- go/arrow/flight/flightsql/client.go | 6 +- .../flightsql/example/sql_batch_reader.go | 19 +- .../flight/flightsql/example/sqlite_server.go | 118 +++++++++--- .../sqlite_tables_schema_batch_reader.go | 2 +- .../flight/flightsql/example/sqlite_test.go | 2 +- go/arrow/flight/flightsql/server_test.go | 179 ++++++++++++++++++ 6 files changed, 297 insertions(+), 29 deletions(-) diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index 080cd039a71c7..077f5295c9f4c 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -42,13 +42,15 @@ func NewClient(addr string, auth flight.ClientAuthHandler, middleware []flight.C if err != nil { return nil, err } - return &Client{cl}, nil + return &Client{cl, memory.DefaultAllocator}, nil } // Client wraps a regular Flight RPC Client to provide the FlightSQL // interface functions and methods. type Client struct { Client flight.Client + + Alloc memory.Allocator } func descForCommand(cmd proto.Message) (*flight.FlightDescriptor, error) { @@ -141,7 +143,7 @@ func (c *Client) DoGet(ctx context.Context, in *flight.Ticket, opts ...grpc.Call return nil, err } - return flight.NewRecordReader(stream) + return flight.NewRecordReader(stream, ipc.WithAllocator(c.Alloc)) } // GetTables requests a list of tables from the server, with the provided diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index e142e55e06878..f7cf4583bea94 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -18,6 +18,7 @@ package example import ( "database/sql" + "reflect" "strings" "sync/atomic" @@ -27,9 +28,8 @@ import ( "github.com/apache/arrow/go/v10/arrow/memory" ) -func getArrowType(dbtype string) arrow.DataType { +func getArrowTypeFromString(dbtype string) arrow.DataType { dbtype = strings.ToLower(dbtype) - if strings.HasPrefix(dbtype, "varchar") { return arrow.BinaryTypes.String } @@ -48,6 +48,19 @@ func getArrowType(dbtype string) arrow.DataType { } } +func getArrowType(c *sql.ColumnType) arrow.DataType { + dbtype := strings.ToLower(c.DatabaseTypeName()) + if dbtype == "" { + switch c.ScanType().Kind() { + case reflect.Int, reflect.Int64, reflect.Uint64: + return arrow.PrimitiveTypes.Int64 + case reflect.Float32, reflect.Float64: + return arrow.PrimitiveTypes.Float64 + } + } + return getArrowTypeFromString(dbtype) +} + // func getColumnMetadata(bldr *flightsql.ColumnMetadataBuilder, c *sql.ColumnType) arrow.Metadata { // defer bldr.Clear() // bldr.IsAutoIncrement(false).IsReadOnly(false) @@ -120,7 +133,7 @@ func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, e for i, c := range cols { fields[i].Name = c.Name() fields[i].Nullable, _ = c.Nullable() - fields[i].Type = getArrowType(c.DatabaseTypeName()) + fields[i].Type = getArrowType(c) switch fields[i].Type.ID() { case arrow.INT64: if fields[i].Nullable { diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go index 47dc80de5c9c1..04d729aca127f 100644 --- a/go/arrow/flight/flightsql/example/sqlite_server.go +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -31,10 +31,12 @@ package example import ( "context" "database/sql" + "database/sql/driver" "fmt" "math/rand" "strings" "sync" + "unsafe" "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/array" @@ -42,6 +44,7 @@ import ( "github.com/apache/arrow/go/v10/arrow/flight/flightsql" "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" "github.com/apache/arrow/go/v10/arrow/memory" + "github.com/apache/arrow/go/v10/arrow/scalar" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" _ "modernc.org/sqlite" @@ -99,8 +102,8 @@ func prepareQueryForGetTables(cmd flightsql.GetTables) string { type SQLiteFlightSQLServer struct { flightsql.BaseServer - db *sql.DB - mem memory.Allocator + db *sql.DB + Alloc memory.Allocator prepared sync.Map } @@ -142,7 +145,7 @@ func (s *SQLiteFlightSQLServer) flightInfoForCommand(desc *flight.FlightDescript return &flight.FlightInfo{ Endpoint: []*flight.FlightEndpoint{{Ticket: &flight.Ticket{Ticket: desc.Cmd}}}, FlightDescriptor: desc, - Schema: flight.SerializeSchema(schema, s.mem), + Schema: flight.SerializeSchema(schema, s.Alloc), TotalRecords: -1, TotalBytes: -1, } @@ -169,7 +172,7 @@ func (s *SQLiteFlightSQLServer) DoGetStatement(ctx context.Context, cmd flightsq return nil, nil, err } - reader, err := NewSqlBatchReader(s.mem, rows) + reader, err := NewSqlBatchReader(s.Alloc, rows) if err != nil { return nil, nil, err } @@ -186,7 +189,7 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoCatalogs(_ context.Context, desc *f func (s *SQLiteFlightSQLServer) DoGetCatalogs(context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { // sqlite doesn't support catalogs, this returns an empty record batch schema := schema_ref.Catalogs - batchBldr := array.NewRecordBuilder(s.mem, schema) + batchBldr := array.NewRecordBuilder(s.Alloc, schema) defer batchBldr.Release() ch := make(chan flight.StreamChunk, 1) @@ -203,7 +206,7 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoSchemas(_ context.Context, cmd flig func (s *SQLiteFlightSQLServer) DoGetDBSchemas(context.Context, flightsql.GetDBSchemas) (*arrow.Schema, <-chan flight.StreamChunk, error) { // sqlite doesn't support schemas, this returns an empty record batch schema := schema_ref.DBSchemas - batchBldr := array.NewRecordBuilder(s.mem, schema) + batchBldr := array.NewRecordBuilder(s.Alloc, schema) defer batchBldr.Release() ch := make(chan flight.StreamChunk, 1) @@ -231,14 +234,14 @@ func (s *SQLiteFlightSQLServer) DoGetTables(ctx context.Context, cmd flightsql.G var rdr array.RecordReader - rdr, err = NewSqlBatchReaderWithSchema(s.mem, schema_ref.Tables, rows) + rdr, err = NewSqlBatchReaderWithSchema(s.Alloc, schema_ref.Tables, rows) if err != nil { return nil, nil, err } ch := make(chan flight.StreamChunk, 2) if cmd.GetIncludeSchema() { - rdr, err = NewSqliteTablesSchemaBatchReader(ctx, s.mem, rdr, s.db, query) + rdr, err = NewSqliteTablesSchemaBatchReader(ctx, s.Alloc, rdr, s.db, query) if err != nil { return nil, nil, err } @@ -263,7 +266,7 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req } handle := genRandomString() - s.prepared.Store(handle, stmt) + s.prepared.Store(string(handle), stmt) result.Handle = handle // no way to get the dataset or parameter schemas from sql.DB @@ -272,8 +275,8 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, request flightsql.ActionClosePreparedStatementRequest) error { handle := request.GetPreparedStatementHandle() - if val, loaded := s.prepared.LoadAndDelete(handle); loaded { - stmt := val.(sql.Stmt) + if val, loaded := s.prepared.LoadAndDelete(string(handle)); loaded { + stmt := val.(*sql.Stmt) return stmt.Close() } @@ -281,7 +284,7 @@ func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, requ } func (s *SQLiteFlightSQLServer) GetFlightInfoPreparedStatement(_ context.Context, cmd flightsql.PreparedStatementQuery, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { - _, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) + _, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) if !ok { return nil, status.Error(codes.InvalidArgument, "prepared statement not found") } @@ -295,18 +298,18 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoPreparedStatement(_ context.Context } func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd flightsql.PreparedStatementQuery) (*arrow.Schema, <-chan flight.StreamChunk, error) { - val, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) + val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) if !ok { return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(sql.Stmt) + stmt := val.(*sql.Stmt) rows, err := stmt.QueryContext(ctx) if err != nil { return nil, nil, err } - rdr, err := NewSqlBatchReader(s.mem, rows) + rdr, err := NewSqlBatchReader(s.Alloc, rows) if err != nil { return nil, nil, err } @@ -316,12 +319,83 @@ func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd return rdr.Schema(), ch, nil } -// func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd flightsql.PreparedStatementUpdate, rdr flight.MessageReader) (int64, error) { -// val, ok := s.prepared.Load(cmd.GetPreparedStatementHandle()) -// if !ok { -// return 0, status.Error(codes.InvalidArgument, "prepared statement not found") -// } +type sqlScalar struct { + v scalar.Scalar +} + +func (s sqlScalar) Value() (driver.Value, error) { + switch v := s.v.(type) { + case *scalar.Int64: + return v.Value, nil + case *scalar.Float32: + return v.Value, nil + case *scalar.Float64: + return v.Value, nil + case *scalar.String: + return *(*string)(unsafe.Pointer(&v.Value.Bytes()[0])), nil + case *scalar.Binary: + return v.Value.Bytes(), nil + default: + return nil, fmt.Errorf("unsupported data type: %s", s.v.DataType()) + } +} + +func getParamsForStatement(rdr flight.MessageReader) (params []sqlScalar, err error) { + for rdr.Next() { + rec := rdr.Record() + + nrows := int(rec.NumRows()) + ncols := int(rec.NumCols()) + + if len(params) < int(ncols) { + params = make([]sqlScalar, ncols) + } -// stmt := val.(sql.Stmt) + for i := 0; i < nrows; i++ { + for c := 0; c < ncols; c++ { + col := rec.Column(c) + sc, err := scalar.GetScalar(col, i) + if err != nil { + return nil, err + } + + if params[c].v != nil { + if r, ok := params[c].v.(scalar.Releasable); ok { + r.Release() + } + } + params[c].v = sc.(*scalar.DenseUnion).Value + } + } + } -// } + return params, rdr.Err() +} + +func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd flightsql.PreparedStatementUpdate, rdr flight.MessageReader) (int64, error) { + val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) + if !ok { + return 0, status.Error(codes.InvalidArgument, "prepared statement not found") + } + + stmt := val.(*sql.Stmt) + params, err := getParamsForStatement(rdr) + if err != nil { + return 0, status.Errorf(codes.Internal, "error gathering parameters for prepared statement: %s", err.Error()) + } + + args := make([]interface{}, len(params)) + for i := range params { + if r, ok := params[i].v.(scalar.Releasable); ok { + defer r.Release() + } + args[i] = params[i] + } + + result, err := stmt.ExecContext(ctx, args...) + if err != nil { + return 0, err + } + + return result.RowsAffected() +} diff --git a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go index 90d9bde94e73c..40c88f5100b5c 100644 --- a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go @@ -175,7 +175,7 @@ func (s *SqliteTablesSchemaBatchReader) Next() bool { columnFields = append(columnFields, arrow.Field{ Name: name, - Type: getArrowType(typ), + Type: getArrowTypeFromString(typ), Nullable: nn == 1, Metadata: getColumnMetadata(bldr, getSqlTypeFromTypeName(typ), tableName), }) diff --git a/go/arrow/flight/flightsql/example/sqlite_test.go b/go/arrow/flight/flightsql/example/sqlite_test.go index 52c107ff83bb3..b8e9be831ffbd 100644 --- a/go/arrow/flight/flightsql/example/sqlite_test.go +++ b/go/arrow/flight/flightsql/example/sqlite_test.go @@ -43,7 +43,7 @@ func TestServer(t *testing.T) { srv, err := NewSQLiteFlightSQLServer() assert.NoError(t, err) assert.NotNil(t, srv) - srv.mem = mem + srv.Alloc = mem s.RegisterFlightService(flightsql.NewFlightServer(srv)) s.Init("localhost:0") diff --git a/go/arrow/flight/flightsql/server_test.go b/go/arrow/flight/flightsql/server_test.go index ece7754bbbdb2..7db3f55c7fe23 100644 --- a/go/arrow/flight/flightsql/server_test.go +++ b/go/arrow/flight/flightsql/server_test.go @@ -18,13 +18,19 @@ package flightsql_test import ( "context" + "os" "strings" "testing" + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" "github.com/apache/arrow/go/v10/arrow/flight" "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/example" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" "github.com/apache/arrow/go/v10/arrow/memory" + "github.com/apache/arrow/go/v10/arrow/scalar" "github.com/stretchr/testify/suite" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -210,3 +216,176 @@ func (s *UnimplementedFlightSqlServerSuite) TestDoAction() { func TestBaseServer(t *testing.T) { suite.Run(t, new(UnimplementedFlightSqlServerSuite)) } + +type FlightSqliteServerSuite struct { + suite.Suite + + srv *example.SQLiteFlightSQLServer + s flight.Server + cl *flightsql.Client + + mem *memory.CheckedAllocator +} + +func (s *FlightSqliteServerSuite) SetupTest() { + var err error + s.mem = memory.NewCheckedAllocator(memory.DefaultAllocator) + s.s = flight.NewServerWithMiddleware(nil) + s.srv, err = example.NewSQLiteFlightSQLServer() + s.Require().NoError(err) + s.srv.Alloc = s.mem + + s.s.RegisterFlightService(flightsql.NewFlightServer(s.srv)) + s.s.Init("localhost:0") + s.s.SetShutdownOnSignals(os.Interrupt, os.Kill) + go s.s.Serve() + s.cl, err = flightsql.NewClient(s.s.Addr().String(), nil, nil, dialOpts...) + s.Require().NoError(err) + s.Require().NotNil(s.cl) + s.cl.Alloc = s.mem +} + +func (s *FlightSqliteServerSuite) TearDownTest() { + s.Require().NoError(s.cl.Close()) + s.s.Shutdown() + s.srv = nil + s.mem.AssertSize(s.T(), 0) +} + +func (s *FlightSqliteServerSuite) fromJSON(dt arrow.DataType, json string) arrow.Array { + arr, _, _ := array.FromJSON(s.mem, dt, strings.NewReader(json)) + return arr +} + +func (s *FlightSqliteServerSuite) execCountQuery(query string) int64 { + info, err := s.cl.Execute(context.Background(), query) + s.NoError(err) + + rdr, err := s.cl.DoGet(context.Background(), info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + rec, err := rdr.Read() + s.NoError(err) + return rec.Column(0).(*array.Int64).Value(0) +} + +func (s *FlightSqliteServerSuite) TestCommandStatementQuery() { + ctx := context.Background() + info, err := s.cl.Execute(ctx, "SELECT * FROM intTable") + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + }, nil) + + s.Truef(expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) + + idarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) + defer idarr.Release() + keyarr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) + defer keyarr.Release() + valarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) + defer valarr.Release() + foreignarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) + defer foreignarr.Release() + + expectedRec := array.NewRecord(expectedSchema, []arrow.Array{idarr, keyarr, valarr, foreignarr}, 4) + defer expectedRec.Release() + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTables() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{}) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer catalogName.Release() + schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer schemaName.Release() + + tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) + defer tableName.Release() + + tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) + defer tableType.Release() + + expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableFilter() { + +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdateWithParams() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', ?)") + s.NoError(err) + defer stmt.Close(ctx) + + typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[2]") + offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") + strArray := s.fromJSON(arrow.BinaryTypes.String, "[]") + bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") + bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[999]") + dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") + paramArr, err := array.NewDenseUnionFromArraysWithFields(typeIDs, + offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, + []string{"string", "bytes", "bigint", "double"}) + s.NoError(err) + batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ + {Name: "parameter_1", Type: paramArr.DataType()}}, nil), + []arrow.Array{paramArr}, 1) + defer func() { + typeIDs.Release() + offsets.Release() + strArray.Release() + bytesArr.Release() + bigintArr.Release() + dblArr.Release() + paramArr.Release() + batch.Release() + }() + + stmt.SetParameters(batch) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + n, err := stmt.ExecuteUpdate(context.Background()) + s.NoError(err) + s.EqualValues(1, n) + s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + n, err = s.cl.ExecuteUpdate(context.Background(), "DELETE FROM intTable WHERE keyName = 'new_value'") + s.NoError(err) + s.EqualValues(1, n) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) +} + +func TestSqliteServer(t *testing.T) { + suite.Run(t, new(FlightSqliteServerSuite)) +} From 76d0864efbcfb77b6049356125f20ed638b402ec Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 12 Aug 2022 15:47:45 -0400 Subject: [PATCH 11/23] more updates and handling --- go/arrow/array/booleanbuilder.go | 7 + go/arrow/array/builder.go | 4 +- go/arrow/array/list.go | 44 ++- .../flightsql/example/sql_batch_reader.go | 4 + .../flight/flightsql/example/sql_statement.go | 36 -- .../flight/flightsql/example/sqlite_server.go | 145 ++++--- .../sqlite_tables_schema_batch_reader.go | 2 +- .../flight/flightsql/example/type_info.go | 115 ++++++ go/arrow/flight/flightsql/server_test.go | 372 +++++++++++++++++- 9 files changed, 627 insertions(+), 102 deletions(-) delete mode 100644 go/arrow/flight/flightsql/example/sql_statement.go create mode 100644 go/arrow/flight/flightsql/example/type_info.go diff --git a/go/arrow/array/booleanbuilder.go b/go/arrow/array/booleanbuilder.go index 83e1fc6d31f2e..760d755314a7a 100644 --- a/go/arrow/array/booleanbuilder.go +++ b/go/arrow/array/booleanbuilder.go @@ -187,6 +187,12 @@ func (b *BooleanBuilder) unmarshalOne(dec *json.Decoder) error { return err } b.Append(val) + case json.Number: + val, err := strconv.ParseBool(v.String()) + if err != nil { + return err + } + b.Append(val) case nil: b.AppendNull() default: @@ -210,6 +216,7 @@ func (b *BooleanBuilder) unmarshal(dec *json.Decoder) error { func (b *BooleanBuilder) UnmarshalJSON(data []byte) error { dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() t, err := dec.Token() if err != nil { return err diff --git a/go/arrow/array/builder.go b/go/arrow/array/builder.go index 4733ba9bbee7b..6a2146c080c40 100644 --- a/go/arrow/array/builder.go +++ b/go/arrow/array/builder.go @@ -304,7 +304,7 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder { } case arrow.LIST: typ := dtype.(*arrow.ListType) - return NewListBuilder(mem, typ.Elem()) + return NewListBuilderWithField(mem, typ.ElemField()) case arrow.STRUCT: typ := dtype.(*arrow.StructType) return NewStructBuilder(mem, typ) @@ -319,7 +319,7 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder { return NewDictionaryBuilder(mem, typ) case arrow.LARGE_LIST: typ := dtype.(*arrow.LargeListType) - return NewLargeListBuilder(mem, typ.Elem()) + return NewLargeListBuilderWithField(mem, typ.ElemField()) case arrow.MAP: typ := dtype.(*arrow.MapType) return NewMapBuilder(mem, typ.KeyType(), typ.ItemType(), typ.KeysSorted) diff --git a/go/arrow/array/list.go b/go/arrow/array/list.go index a603f7f7ada66..a912839e45cda 100644 --- a/go/arrow/array/list.go +++ b/go/arrow/array/list.go @@ -321,12 +321,31 @@ func NewListBuilder(mem memory.Allocator, etype arrow.DataType) *ListBuilder { } } +// NewListBuilderWithField takes a field to use for the child rather than just +// a datatype to allow for more customization. +func NewListBuilderWithField(mem memory.Allocator, field arrow.Field) *ListBuilder { + offsetBldr := NewInt32Builder(mem) + return &ListBuilder{ + baseListBuilder{ + builder: builder{refCount: 1, mem: mem}, + values: NewBuilder(mem, field.Type), + offsets: offsetBldr, + dt: arrow.ListOfField(field), + appendOffsetVal: func(o int) { offsetBldr.Append(int32(o)) }, + }, + } +} + func (b *baseListBuilder) Type() arrow.DataType { - switch b.dt.ID() { - case arrow.LIST: - return arrow.ListOf(b.values.Type()) - case arrow.LARGE_LIST: - return arrow.LargeListOf(b.values.Type()) + switch dt := b.dt.(type) { + case *arrow.ListType: + f := dt.ElemField() + f.Type = b.values.Type() + return arrow.ListOfField(f) + case *arrow.LargeListType: + f := dt.ElemField() + f.Type = b.values.Type() + return arrow.LargeListOfField(f) } return nil } @@ -346,6 +365,21 @@ func NewLargeListBuilder(mem memory.Allocator, etype arrow.DataType) *LargeListB } } +// NewLargeListBuilderWithField takes a field rather than just an element type +// to allow for more customization of the final type of the LargeList Array +func NewLargeListBuilderWithField(mem memory.Allocator, field arrow.Field) *ListBuilder { + offsetBldr := NewInt64Builder(mem) + return &ListBuilder{ + baseListBuilder{ + builder: builder{refCount: 1, mem: mem}, + values: NewBuilder(mem, field.Type), + offsets: offsetBldr, + dt: arrow.LargeListOfField(field), + appendOffsetVal: func(o int) { offsetBldr.Append(int64(o)) }, + }, + } +} + // Release decreases the reference count by 1. // When the reference count goes to zero, the memory is freed. func (b *baseListBuilder) Release() { diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index f7cf4583bea94..4734a30488ac1 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -24,6 +24,7 @@ import ( "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" "github.com/apache/arrow/go/v10/arrow/internal/debug" "github.com/apache/arrow/go/v10/arrow/memory" ) @@ -122,6 +123,8 @@ func NewSqlBatchReaderWithSchema(mem memory.Allocator, schema *arrow.Schema, row } func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, error) { + bldr := flightsql.NewColumnMetadataBuilder() + cols, err := rows.ColumnTypes() if err != nil { rows.Close() @@ -134,6 +137,7 @@ func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, e fields[i].Name = c.Name() fields[i].Nullable, _ = c.Nullable() fields[i].Type = getArrowType(c) + fields[i].Metadata = getColumnMetadata(bldr, getSqlTypeFromTypeName(c.DatabaseTypeName()), "") switch fields[i].Type.ID() { case arrow.INT64: if fields[i].Nullable { diff --git a/go/arrow/flight/flightsql/example/sql_statement.go b/go/arrow/flight/flightsql/example/sql_statement.go deleted file mode 100644 index 6ae75a55d9122..0000000000000 --- a/go/arrow/flight/flightsql/example/sql_statement.go +++ /dev/null @@ -1,36 +0,0 @@ -// 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. - -package example - -import ( - "database/sql" -) - -// func GetColumnMetadata(coltype int, table string) flightsql.ColumnMetadata {} - -type Statement struct { - stmt *sql.Stmt - db *sql.DB -} - -func NewStatement(db *sql.DB, query string) (*Statement, error) { - stmt, err := db.Prepare(query) - if err != nil { - return nil, err - } - return &Statement{stmt, db}, nil -} diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go index 04d729aca127f..f89c2835fc3dd 100644 --- a/go/arrow/flight/flightsql/example/sqlite_server.go +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -31,12 +31,10 @@ package example import ( "context" "database/sql" - "database/sql/driver" "fmt" "math/rand" "strings" "sync" - "unsafe" "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/array" @@ -100,6 +98,11 @@ func prepareQueryForGetTables(cmd flightsql.GetTables) string { return b.String() } +type Statement struct { + stmt *sql.Stmt + params []interface{} +} + type SQLiteFlightSQLServer struct { flightsql.BaseServer db *sql.DB @@ -189,11 +192,8 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoCatalogs(_ context.Context, desc *f func (s *SQLiteFlightSQLServer) DoGetCatalogs(context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { // sqlite doesn't support catalogs, this returns an empty record batch schema := schema_ref.Catalogs - batchBldr := array.NewRecordBuilder(s.Alloc, schema) - defer batchBldr.Release() - ch := make(chan flight.StreamChunk, 1) - ch <- flight.StreamChunk{Data: batchBldr.NewRecord()} + ch := make(chan flight.StreamChunk) close(ch) return schema, ch, nil @@ -206,11 +206,8 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoSchemas(_ context.Context, cmd flig func (s *SQLiteFlightSQLServer) DoGetDBSchemas(context.Context, flightsql.GetDBSchemas) (*arrow.Schema, <-chan flight.StreamChunk, error) { // sqlite doesn't support schemas, this returns an empty record batch schema := schema_ref.DBSchemas - batchBldr := array.NewRecordBuilder(s.Alloc, schema) - defer batchBldr.Release() - ch := make(chan flight.StreamChunk, 1) - ch <- flight.StreamChunk{Data: batchBldr.NewRecord()} + ch := make(chan flight.StreamChunk) close(ch) return schema, ch, nil @@ -251,6 +248,46 @@ func (s *SQLiteFlightSQLServer) DoGetTables(ctx context.Context, cmd flightsql.G return rdr.Schema(), ch, nil } +func (s *SQLiteFlightSQLServer) GetFlightInfoXdbcTypeInfo(_ context.Context, _ flightsql.GetXdbcTypeInfo, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.XdbcTypeInfo), nil +} + +func (s *SQLiteFlightSQLServer) DoGetXdbcTypeInfo(_ context.Context, cmd flightsql.GetXdbcTypeInfo) (*arrow.Schema, <-chan flight.StreamChunk, error) { + var batch arrow.Record + if cmd.GetDataType() == nil { + batch = GetTypeInfoResult(s.Alloc) + } else { + batch = GetFilteredTypeInfoResult(s.Alloc, *cmd.GetDataType()) + } + + ch := make(chan flight.StreamChunk, 1) + ch <- flight.StreamChunk{Data: batch} + close(ch) + return batch.Schema(), ch, nil +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoTableTypes(_ context.Context, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.TableTypes), nil +} + +func (s *SQLiteFlightSQLServer) DoGetTableTypes(ctx context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { + query := "SELECT DISTINCT type AS table_type FROM sqlite_master" + + rows, err := s.db.QueryContext(ctx, query) + if err != nil { + return nil, nil, err + } + + reader, err := NewSqlBatchReaderWithSchema(s.Alloc, schema_ref.TableTypes, rows) + if err != nil { + return nil, nil, err + } + + ch := make(chan flight.StreamChunk) + go flight.StreamChunksFromReader(reader, ch) + return reader.schema, ch, nil +} + func (s *SQLiteFlightSQLServer) DoPutCommandStatementUpdate(ctx context.Context, cmd flightsql.StatementUpdate) (int64, error) { res, err := s.db.ExecContext(ctx, cmd.GetQuery()) if err != nil { @@ -266,7 +303,7 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req } handle := genRandomString() - s.prepared.Store(string(handle), stmt) + s.prepared.Store(string(handle), Statement{stmt: stmt}) result.Handle = handle // no way to get the dataset or parameter schemas from sql.DB @@ -276,8 +313,8 @@ func (s *SQLiteFlightSQLServer) CreatePreparedStatement(ctx context.Context, req func (s *SQLiteFlightSQLServer) ClosePreparedStatement(ctx context.Context, request flightsql.ActionClosePreparedStatementRequest) error { handle := request.GetPreparedStatementHandle() if val, loaded := s.prepared.LoadAndDelete(string(handle)); loaded { - stmt := val.(*sql.Stmt) - return stmt.Close() + stmt := val.(Statement) + return stmt.stmt.Close() } return status.Error(codes.InvalidArgument, "prepared statement not found") @@ -303,8 +340,8 @@ func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd return nil, nil, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(*sql.Stmt) - rows, err := stmt.QueryContext(ctx) + stmt := val.(Statement) + rows, err := stmt.stmt.QueryContext(ctx, stmt.params...) if err != nil { return nil, nil, err } @@ -319,28 +356,7 @@ func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd return rdr.Schema(), ch, nil } -type sqlScalar struct { - v scalar.Scalar -} - -func (s sqlScalar) Value() (driver.Value, error) { - switch v := s.v.(type) { - case *scalar.Int64: - return v.Value, nil - case *scalar.Float32: - return v.Value, nil - case *scalar.Float64: - return v.Value, nil - case *scalar.String: - return *(*string)(unsafe.Pointer(&v.Value.Bytes()[0])), nil - case *scalar.Binary: - return v.Value.Bytes(), nil - default: - return nil, fmt.Errorf("unsupported data type: %s", s.v.DataType()) - } -} - -func getParamsForStatement(rdr flight.MessageReader) (params []sqlScalar, err error) { +func getParamsForStatement(rdr flight.MessageReader) (params []interface{}, err error) { for rdr.Next() { rec := rdr.Record() @@ -348,7 +364,7 @@ func getParamsForStatement(rdr flight.MessageReader) (params []sqlScalar, err er ncols := int(rec.NumCols()) if len(params) < int(ncols) { - params = make([]sqlScalar, ncols) + params = make([]interface{}, ncols) } for i := 0; i < nrows; i++ { @@ -358,13 +374,24 @@ func getParamsForStatement(rdr flight.MessageReader) (params []sqlScalar, err er if err != nil { return nil, err } + if r, ok := sc.(scalar.Releasable); ok { + r.Release() + } - if params[c].v != nil { - if r, ok := params[c].v.(scalar.Releasable); ok { - r.Release() - } + switch v := sc.(*scalar.DenseUnion).Value.(type) { + case *scalar.Int64: + params[c] = v.Value + case *scalar.Float32: + params[c] = v.Value + case *scalar.Float64: + params[c] = v.Value + case *scalar.String: + params[c] = string(v.Value.Bytes()) + case *scalar.Binary: + params[c] = v.Value.Bytes() + default: + return nil, fmt.Errorf("unsupported type: %s", v) } - params[c].v = sc.(*scalar.DenseUnion).Value } } } @@ -372,27 +399,37 @@ func getParamsForStatement(rdr flight.MessageReader) (params []sqlScalar, err er return params, rdr.Err() } +func (s *SQLiteFlightSQLServer) DoPutPreparedStatementQuery(_ context.Context, cmd flightsql.PreparedStatementQuery, rdr flight.MessageReader, _ flight.MetadataWriter) error { + val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) + if !ok { + return status.Error(codes.InvalidArgument, "prepared statement not found") + } + + stmt := val.(Statement) + args, err := getParamsForStatement(rdr) + if err != nil { + return status.Errorf(codes.Internal, "error gathering parameters for prepared statement query: %s", err.Error()) + } + + stmt.params = args + s.prepared.Store(string(cmd.GetPreparedStatementHandle()), stmt) + return nil +} + func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd flightsql.PreparedStatementUpdate, rdr flight.MessageReader) (int64, error) { val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) if !ok { return 0, status.Error(codes.InvalidArgument, "prepared statement not found") } - stmt := val.(*sql.Stmt) - params, err := getParamsForStatement(rdr) + stmt := val.(Statement) + args, err := getParamsForStatement(rdr) if err != nil { return 0, status.Errorf(codes.Internal, "error gathering parameters for prepared statement: %s", err.Error()) } - args := make([]interface{}, len(params)) - for i := range params { - if r, ok := params[i].v.(scalar.Releasable); ok { - defer r.Release() - } - args[i] = params[i] - } - - result, err := stmt.ExecContext(ctx, args...) + stmt.params = args + result, err := stmt.stmt.ExecContext(ctx, args...) if err != nil { return 0, err } diff --git a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go index 40c88f5100b5c..005b7e5317e0b 100644 --- a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go @@ -128,7 +128,7 @@ func getPrecisionFromCol(sqltype int) int { func getColumnMetadata(bldr *flightsql.ColumnMetadataBuilder, sqltype int, table string) arrow.Metadata { defer bldr.Clear() - bldr.Scale(15).IsReadOnly(false) + bldr.Scale(15).IsReadOnly(false).IsAutoIncrement(false) if table != "" { bldr.TableName(table) } diff --git a/go/arrow/flight/flightsql/example/type_info.go b/go/arrow/flight/flightsql/example/type_info.go new file mode 100644 index 0000000000000..9e301bddf9e57 --- /dev/null +++ b/go/arrow/flight/flightsql/example/type_info.go @@ -0,0 +1,115 @@ +// 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. + +package example + +import ( + "strings" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" + "github.com/apache/arrow/go/v10/arrow/memory" +) + +func GetTypeInfoResult(mem memory.Allocator) arrow.Record { + typeNames, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, + strings.NewReader(`["bit", "tinyint", "bigint", "longvarbinary", + "varbinary", "text", "longvarchar", "char", + "integer", "smallint", "float", "double", + "numeric", "varchar", "date", "time", "timestamp"]`)) + defer typeNames.Release() + + dataType, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, + strings.NewReader(`[-7, -6, -5, -4, -3, -1, -1, 1, 4, 5, 6, 8, 8, 12, 91, 92, 93]`)) + defer dataType.Release() + + columnSize, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, + strings.NewReader(`[1, 3, 19, 65536, 255, 65536, 65536, 255, 9, 5, 7, 15, 15, 255, 10, 8, 32]`)) + defer columnSize.Release() + + literalPrefix, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, + strings.NewReader(`[null, null, null, null, null, "'", "'", "'", null, null, null, null, null, "'" ,"'", "'", "'"]`)) + defer literalPrefix.Release() + + literalSuffix, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, + strings.NewReader(`[null, null, null, null, null, "'", "'", "'", null, null, null, null, null, "'" ,"'", "'", "'"]`)) + defer literalSuffix.Release() + + createParams, _, _ := array.FromJSON(mem, arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.BinaryTypes.String, Nullable: false}), + strings.NewReader(`[[], [], [], [], [], ["length"], ["length"], ["length"], [], [], [], [], [], ["length"], [], [], []]`)) + defer createParams.Release() + + nullable, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, + strings.NewReader(`[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]`)) + defer nullable.Release() + + // reference for creating a boolean() array with only zeros + zeroBoolArray, _, err := array.FromJSON(mem, arrow.FixedWidthTypes.Boolean, + strings.NewReader(`[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]`), array.WithUseNumber()) + if err != nil { + panic(err) + } + defer zeroBoolArray.Release() + caseSensitive := zeroBoolArray + + searchable, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, + strings.NewReader(`[3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]`)) + defer searchable.Release() + + unsignedAttribute := zeroBoolArray + fixedPrecScale := zeroBoolArray + autoUniqueVal := zeroBoolArray + + localTypeName := typeNames + + zeroIntArray, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int32, + strings.NewReader(`[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]`)) + defer zeroIntArray.Release() + + minimalScale := zeroIntArray + maximumScale := zeroIntArray + sqlDataType := dataType + sqlDateTimeSub := zeroIntArray + numPrecRadix := zeroIntArray + intervalPrecision := zeroIntArray + + return array.NewRecord(schema_ref.XdbcTypeInfo, []arrow.Array{ + typeNames, dataType, columnSize, literalPrefix, literalSuffix, + createParams, nullable, caseSensitive, searchable, unsignedAttribute, + fixedPrecScale, autoUniqueVal, localTypeName, minimalScale, maximumScale, + sqlDataType, sqlDateTimeSub, numPrecRadix, intervalPrecision}, 17) +} + +func GetFilteredTypeInfoResult(mem memory.Allocator, filter int32) arrow.Record { + batch := GetTypeInfoResult(mem) + defer batch.Release() + + dataTypeVector := []int32{-7, -6, -5, -4, -3, -1, -1, 1, 4, 5, 6, 8, 8, 12, 91, 92, 93} + start, end := -1, -1 + for i, v := range dataTypeVector { + if filter == v { + if start == -1 { + start = i + } + } else if start != -1 && end == -1 { + end = i + break + } + } + + return batch.NewSlice(int64(start), int64(end)) +} diff --git a/go/arrow/flight/flightsql/server_test.go b/go/arrow/flight/flightsql/server_test.go index 7db3f55c7fe23..e488c802be6dd 100644 --- a/go/arrow/flight/flightsql/server_test.go +++ b/go/arrow/flight/flightsql/server_test.go @@ -38,6 +38,7 @@ import ( "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/anypb" + sqlite3 "modernc.org/sqlite/lib" ) var dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} @@ -227,6 +228,24 @@ type FlightSqliteServerSuite struct { mem *memory.CheckedAllocator } +func (s *FlightSqliteServerSuite) getColMetadata(colType int, table string) arrow.Metadata { + bldr := flightsql.NewColumnMetadataBuilder() + bldr.Scale(15).IsReadOnly(false).IsAutoIncrement(false) + if table != "" { + bldr.TableName(table) + } + switch colType { + case sqlite3.SQLITE_TEXT, sqlite3.SQLITE_BLOB: + case sqlite3.SQLITE_INTEGER: + bldr.Precision(10) + case sqlite3.SQLITE_FLOAT: + bldr.Precision(15) + default: + bldr.Precision(0) + } + return bldr.Metadata() +} + func (s *FlightSqliteServerSuite) SetupTest() { var err error s.mem = memory.NewCheckedAllocator(memory.DefaultAllocator) @@ -283,10 +302,10 @@ func (s *FlightSqliteServerSuite) TestCommandStatementQuery() { s.NotNil(rec) expectedSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "keyName", Type: arrow.BinaryTypes.String, Nullable: true}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, }, nil) s.Truef(expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) @@ -341,7 +360,335 @@ func (s *FlightSqliteServerSuite) TestCommandGetTables() { } func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableNameFilterPattern: proto.String("int%"), + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) + tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + expected := array.NewRecord(schema_ref.Tables, []arrow.Array{catalog, schema, table, tabletype}, 1) + defer func() { + catalog.Release() + schema.Release() + table.Release() + tabletype.Release() + expected.Release() + }() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + s.NoError(rdr.Err()) + + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableTypesFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableTypes: []string{"index"}, + }) + s.NoError(err) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(schema_ref.Tables.Equal(rdr.Schema()), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithExistingTableTypeFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableTypes: []string{"table"}, + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer catalogName.Release() + schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer schemaName.Release() + + tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) + defer tableName.Release() + + tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) + defer tableType.Release() + + expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithIncludedSchemas() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableNameFilterPattern: proto.String("int%"), + IncludeSchema: true, + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) + tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + + dbTableName := "intTable" + + tableSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + {Name: "keyName", Type: arrow.BinaryTypes.String, + Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, dbTableName)}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + }, nil) + schemaBuf := flight.SerializeSchema(tableSchema, s.mem) + binaryBldr := array.NewBinaryBuilder(s.mem, arrow.BinaryTypes.Binary) + binaryBldr.Append(schemaBuf) + schemaCol := binaryBldr.NewArray() + + expected := array.NewRecord(schema_ref.TablesWithIncludedSchema, []arrow.Array{catalog, schema, table, tabletype, schemaCol}, 1) + defer func() { + catalog.Release() + schema.Release() + table.Release() + tabletype.Release() + binaryBldr.Release() + schemaCol.Release() + expected.Release() + }() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + s.NoError(rdr.Err()) + + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTypeInfo() { + ctx := context.Background() + info, err := s.cl.GetXdbcTypeInfo(ctx, nil) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := example.GetTypeInfoResult(s.mem) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTypeInfoFiltered() { + ctx := context.Background() + info, err := s.cl.GetXdbcTypeInfo(ctx, proto.Int32(-4)) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := example.GetFilteredTypeInfoResult(s.mem, -4) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetCatalogs() { + ctx := context.Background() + info, err := s.cl.GetCatalogs(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Schema().Equal(schema_ref.Catalogs), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetDbSchemas() { + ctx := context.Background() + info, err := s.cl.GetDBSchemas(ctx, &flightsql.GetDBSchemasOpts{}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Schema().Equal(schema_ref.DBSchemas), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTableTypes() { + ctx := context.Background() + info, err := s.cl.GetTableTypes(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + defer expected.Release() + expectedRec := array.NewRecord(schema_ref.TableTypes, []arrow.Array{expected}, 1) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandStatementUpdate() { + ctx := context.Background() + result, err := s.cl.ExecuteUpdate(ctx, `INSERT INTO intTable (keyName, value) VALUES + ('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)`) + s.NoError(err) + s.EqualValues(3, result) + + result, err = s.cl.ExecuteUpdate(ctx, `UPDATE intTable SET keyName = 'KEYNAME1' + WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'`) + s.NoError(err) + s.EqualValues(2, result) + + result, err = s.cl.ExecuteUpdate(ctx, `DELETE FROM intTable WHERE keyName = 'KEYNAME1'`) + s.NoError(err) + s.EqualValues(3, result) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQuery() { + ctx := context.Background() + prep, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable") + s.NoError(err) + defer prep.Close(ctx) + + info, err := prep.Execute(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) + + idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) + defer idArr.Release() + keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) + defer keyNameArr.Release() + valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) + defer valueArr.Release() + foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) + defer foreignIdArr.Release() + + expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 4) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQueryWithParams() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable WHERE keyName LIKE ?") + s.NoError(err) + defer stmt.Close(ctx) + + typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[0]") + offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") + strArray := s.fromJSON(arrow.BinaryTypes.String, `["%one"]`) + bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") + bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[]") + dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") + paramArr, _ := array.NewDenseUnionFromArraysWithFields(typeIDs, + offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, + []string{"string", "bytes", "bigint", "double"}) + batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ + {Name: "parameter_1", Type: paramArr.DataType()}}, nil), + []arrow.Array{paramArr}, 1) + defer func() { + typeIDs.Release() + offsets.Release() + strArray.Release() + bytesArr.Release() + bigintArr.Release() + dblArr.Release() + paramArr.Release() + batch.Release() + }() + stmt.SetParameters(batch) + info, err := stmt.Execute(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) + + idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 3]`) + defer idArr.Release() + keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "negative one"]`) + defer keyNameArr.Release() + valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, -1]`) + defer valueArr.Release() + foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1]`) + defer foreignIdArr.Release() + + expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 2) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) } func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdateWithParams() { @@ -386,6 +733,23 @@ func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdateWithParams() s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) } +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdate() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', 999)") + s.NoError(err) + defer stmt.Close(ctx) + + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + result, err := stmt.ExecuteUpdate(ctx) + s.NoError(err) + s.EqualValues(1, result) + s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + result, err = s.cl.ExecuteUpdate(ctx, "DELETE FROM intTable WHERE keyName = 'new_value'") + s.NoError(err) + s.EqualValues(1, result) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) +} + func TestSqliteServer(t *testing.T) { suite.Run(t, new(FlightSqliteServerSuite)) } From 24529bc5bd3427e07a9711e49d2cc9abf75df0d2 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Fri, 12 Aug 2022 17:43:30 -0400 Subject: [PATCH 12/23] completed sqlite server example --- go/arrow/array/list.go | 13 +- .../flightsql/example/sql_batch_reader.go | 50 ++++- .../flight/flightsql/example/sqlite_info.go | 193 ++++++++++++++++ .../flight/flightsql/example/sqlite_server.go | 179 ++++++++++++--- go/arrow/flight/flightsql/server.go | 4 + go/arrow/flight/flightsql/server_test.go | 208 ++++++++++++++++++ 6 files changed, 599 insertions(+), 48 deletions(-) create mode 100644 go/arrow/flight/flightsql/example/sqlite_info.go diff --git a/go/arrow/array/list.go b/go/arrow/array/list.go index a912839e45cda..5b2e4b3f61f38 100644 --- a/go/arrow/array/list.go +++ b/go/arrow/array/list.go @@ -390,10 +390,10 @@ func (b *baseListBuilder) Release() { b.nullBitmap.Release() b.nullBitmap = nil } + b.values.Release() + b.offsets.Release() } - b.values.Release() - b.offsets.Release() } func (b *baseListBuilder) appendNextOffset() { @@ -488,9 +488,6 @@ func (b *LargeListBuilder) NewArray() arrow.Array { // NewListArray creates a List array from the memory buffers used by the builder and resets the ListBuilder // so it can be used to build a new array. func (b *ListBuilder) NewListArray() (a *List) { - if b.offsets.Len() != b.length+1 { - b.appendNextOffset() - } data := b.newData() a = NewListData(data) data.Release() @@ -500,9 +497,6 @@ func (b *ListBuilder) NewListArray() (a *List) { // NewLargeListArray creates a List array from the memory buffers used by the builder and resets the LargeListBuilder // so it can be used to build a new array. func (b *LargeListBuilder) NewLargeListArray() (a *LargeList) { - if b.offsets.Len() != b.length+1 { - b.appendNextOffset() - } data := b.newData() a = NewLargeListData(data) data.Release() @@ -510,6 +504,9 @@ func (b *LargeListBuilder) NewLargeListArray() (a *LargeList) { } func (b *baseListBuilder) newData() (data *Data) { + if b.offsets.Len() != b.length+1 { + b.appendNextOffset() + } values := b.values.NewArray() defer values.Release() diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index 4734a30488ac1..d56f057a477c3 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -62,16 +62,6 @@ func getArrowType(c *sql.ColumnType) arrow.DataType { return getArrowTypeFromString(dbtype) } -// func getColumnMetadata(bldr *flightsql.ColumnMetadataBuilder, c *sql.ColumnType) arrow.Metadata { -// defer bldr.Clear() -// bldr.IsAutoIncrement(false).IsReadOnly(false) -// prec, scale, ok := c.DecimalSize() -// if ok { -// bldr.Precision(int32(prec)).Scale(int32(scale)) -// } -// return bldr.Metadata() -// } - const maxBatchSize = 1024 type SqlBatchReader struct { @@ -90,6 +80,18 @@ func NewSqlBatchReaderWithSchema(mem memory.Allocator, schema *arrow.Schema, row rowdest := make([]interface{}, len(schema.Fields())) for i, f := range schema.Fields() { switch f.Type.ID() { + case arrow.UINT8: + if f.Nullable { + rowdest[i] = &sql.NullInt16{} + } else { + rowdest[i] = new(uint8) + } + case arrow.INT32: + if f.Nullable { + rowdest[i] = &sql.NullInt32{} + } else { + rowdest[i] = new(int32) + } case arrow.INT64: if f.Nullable { rowdest[i] = &sql.NullInt64{} @@ -139,6 +141,18 @@ func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, e fields[i].Type = getArrowType(c) fields[i].Metadata = getColumnMetadata(bldr, getSqlTypeFromTypeName(c.DatabaseTypeName()), "") switch fields[i].Type.ID() { + case arrow.UINT8: + if fields[i].Nullable { + rowdest[i] = &sql.NullInt16{} + } else { + rowdest[i] = new(uint8) + } + case arrow.INT32: + if fields[i].Nullable { + rowdest[i] = &sql.NullInt32{} + } else { + rowdest[i] = new(int32) + } case arrow.INT64: if fields[i].Nullable { rowdest[i] = &sql.NullInt64{} @@ -212,6 +226,14 @@ func (r *SqlBatchReader) Next() bool { for i, v := range r.rowdest { fb := r.bldr.Field(i) switch v := v.(type) { + case *uint8: + fb.(*array.Uint8Builder).Append(*v) + case *sql.NullInt16: + if !v.Valid { + fb.AppendNull() + } else { + fb.(*array.Uint8Builder).Append(uint8(v.Int16)) + } case *int64: fb.(*array.Int64Builder).Append(*v) case *sql.NullInt64: @@ -220,6 +242,14 @@ func (r *SqlBatchReader) Next() bool { } else { fb.(*array.Int64Builder).Append(v.Int64) } + case *int32: + fb.(*array.Int32Builder).Append(*v) + case *sql.NullInt32: + if !v.Valid { + fb.AppendNull() + } else { + fb.(*array.Int32Builder).Append(v.Int32) + } case *float64: fb.(*array.Float64Builder).Append(*v) case *sql.NullFloat64: diff --git a/go/arrow/flight/flightsql/example/sqlite_info.go b/go/arrow/flight/flightsql/example/sqlite_info.go new file mode 100644 index 0000000000000..bc7f34c56e571 --- /dev/null +++ b/go/arrow/flight/flightsql/example/sqlite_info.go @@ -0,0 +1,193 @@ +// 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. + +package example + +import "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + +func SqlInfoResultMap() flightsql.SqlInfoResultMap { + return flightsql.SqlInfoResultMap{ + uint32(flightsql.SqlInfoFlightSqlServerName): "db_name", + uint32(flightsql.SqlInfoFlightSqlServerVersion): "sqlite 3", + uint32(flightsql.SqlInfoFlightSqlServerArrowVersion): "10.0.0-SNAPSHOT", + uint32(flightsql.SqlInfoFlightSqlServerReadOnly): false, + uint32(flightsql.SqlInfoDDLCatalog): false, + uint32(flightsql.SqlInfoDDLSchema): false, + uint32(flightsql.SqlInfoDDLTable): true, + uint32(flightsql.SqlInfoIdentifierCase): int64(flightsql.SqlCaseSensitivityCaseInsensitive), + uint32(flightsql.SqlInfoIdentifierQuoteChar): `"`, + uint32(flightsql.SqlInfoQuotedIdentifierCase): int64(flightsql.SqlCaseSensitivityCaseInsensitive), + uint32(flightsql.SqlInfoAllTablesAreASelectable): true, + uint32(flightsql.SqlInfoNullOrdering): int64(flightsql.SqlNullOrderingSortAtStart), + uint32(flightsql.SqlInfoKeywords): []string{"ABORT", + "ACTION", + "ADD", + "AFTER", + "ALL", + "ALTER", + "ALWAYS", + "ANALYZE", + "AND", + "AS", + "ASC", + "ATTACH", + "AUTOINCREMENT", + "BEFORE", + "BEGIN", + "BETWEEN", + "BY", + "CASCADE", + "CASE", + "CAST", + "CHECK", + "COLLATE", + "COLUMN", + "COMMIT", + "CONFLICT", + "CONSTRAINT", + "CREATE", + "CROSS", + "CURRENT", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "DATABASE", + "DEFAULT", + "DEFERRABLE", + "DEFERRED", + "DELETE", + "DESC", + "DETACH", + "DISTINCT", + "DO", + "DROP", + "EACH", + "ELSE", + "END", + "ESCAPE", + "EXCEPT", + "EXCLUDE", + "EXCLUSIVE", + "EXISTS", + "EXPLAIN", + "FAIL", + "FILTER", + "FIRST", + "FOLLOWING", + "FOR", + "FOREIGN", + "FROM", + "FULL", + "GENERATED", + "GLOB", + "GROUP", + "GROUPS", + "HAVING", + "IF", + "IGNORE", + "IMMEDIATE", + "IN", + "INDEX", + "INDEXED", + "INITIALLY", + "INNER", + "INSERT", + "INSTEAD", + "INTERSECT", + "INTO", + "IS", + "ISNULL", + "JOIN", + "KEY", + "LAST", + "LEFT", + "LIKE", + "LIMIT", + "MATCH", + "MATERIALIZED", + "NATURAL", + "NO", + "NOT", + "NOTHING", + "NOTNULL", + "NULL", + "NULLS", + "OF", + "OFFSET", + "ON", + "OR", + "ORDER", + "OTHERS", + "OUTER", + "OVER", + "PARTITION", + "PLAN", + "PRAGMA", + "PRECEDING", + "PRIMARY", + "QUERY", + "RAISE", + "RANGE", + "RECURSIVE", + "REFERENCES", + "REGEXP", + "REINDEX", + "RELEASE", + "RENAME", + "REPLACE", + "RESTRICT", + "RETURNING", + "RIGHT", + "ROLLBACK", + "ROW", + "ROWS", + "SAVEPOINT", + "SELECT", + "SET", + "TABLE", + "TEMP", + "TEMPORARY", + "THEN", + "TIES", + "TO", + "TRANSACTION", + "TRIGGER", + "UNBOUNDED", + "UNION", + "UNIQUE", + "UPDATE", + "USING", + "VACUUM", + "VALUES", + "VIEW", + "VIRTUAL", + "WHEN", + "WHERE", + "WINDOW", + "WITH", + "WITHOUT"}, + uint32(flightsql.SqlInfoNumericFunctions): []string{ + "ACOS", "ACOSH", "ASIN", "ASINH", "ATAN", "ATAN2", "ATANH", "CEIL", + "CEILING", "COS", "COSH", "DEGREES", "EXP", "FLOOR", "LN", "LOG", + "LOG10", "LOG2", "MOD", "PI", "POW", "POWER", "RADIANS", + "SIN", "SINH", "SQRT", "TAN", "TANH", "TRUNC"}, + uint32(flightsql.SqlInfoStringFunctions): []string{"SUBSTR", "TRIM", "LTRIM", "RTRIM", "LENGTH", + "REPLACE", "UPPER", "LOWER", "INSTR"}, + uint32(flightsql.SqlInfoSupportsConvert): map[int32][]int32{ + int32(flightsql.SqlConvertBigInt): {int32(flightsql.SqlConvertInteger)}, + }, + } +} diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go index f89c2835fc3dd..40e1c0979988c 100644 --- a/go/arrow/flight/flightsql/example/sqlite_server.go +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -98,6 +98,40 @@ func prepareQueryForGetTables(cmd flightsql.GetTables) string { return b.String() } +func prepareQueryForGetKeys(filter string) string { + return `SELECT * FROM ( + SELECT + NULL AS pk_catalog_name, + NULL AS pk_schema_name, + p."table" AS pk_table_name, + p."to" AS pk_column_name, + NULL AS fk_catalog_name, + NULL AS fk_schema_name, + m.name AS fk_table_name, + p."from" AS fk_column_name, + p.seq AS key_sequence, + NULL AS pk_key_name, + NULL AS fk_key_name, + CASE + WHEN p.on_update = 'CASCADE' THEN 0 + WHEN p.on_update = 'RESTRICT' THEN 1 + WHEN p.on_update = 'SET NULL' THEN 2 + WHEN p.on_update = 'NO ACTION' THEN 3 + WHEN p.on_update = 'SET DEFAULT' THEN 4 + END AS update_rule, + CASE + WHEN p.on_delete = 'CASCADE' THEN 0 + WHEN p.on_delete = 'RESTRICT' THEN 1 + WHEN p.on_delete = 'SET NULL' THEN 2 + WHEN p.on_delete = 'NO ACTION' THEN 3 + WHEN p.on_delete = 'SET DEFAULT' THEN 4 + END AS delete_rule + FROM sqlite_master m + JOIN pragma_foreign_key_list(m.name) p ON m.name != p."table" + WHERE m.type = 'table') WHERE ` + filter + + ` ORDER BY pk_catalog_name, pk_schema_name, pk_table_name, pk_key_name, key_sequence` +} + type Statement struct { stmt *sql.Stmt params []interface{} @@ -105,8 +139,7 @@ type Statement struct { type SQLiteFlightSQLServer struct { flightsql.BaseServer - db *sql.DB - Alloc memory.Allocator + db *sql.DB prepared sync.Map } @@ -141,7 +174,11 @@ func NewSQLiteFlightSQLServer() (*SQLiteFlightSQLServer, error) { if err != nil { return nil, err } - return &SQLiteFlightSQLServer{db: db}, nil + ret := &SQLiteFlightSQLServer{db: db} + for k, v := range SqlInfoResultMap() { + ret.RegisterSqlInfo(flightsql.SqlInfo(k), v) + } + return ret, nil } func (s *SQLiteFlightSQLServer) flightInfoForCommand(desc *flight.FlightDescriptor, schema *arrow.Schema) *flight.FlightInfo { @@ -170,19 +207,7 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoStatement(ctx context.Context, cmd } func (s *SQLiteFlightSQLServer) DoGetStatement(ctx context.Context, cmd flightsql.StatementQueryTicket) (*arrow.Schema, <-chan flight.StreamChunk, error) { - rows, err := s.db.QueryContext(ctx, string(cmd.GetStatementHandle())) - if err != nil { - return nil, nil, err - } - - reader, err := NewSqlBatchReader(s.Alloc, rows) - if err != nil { - return nil, nil, err - } - - ch := make(chan flight.StreamChunk) - go flight.StreamChunksFromReader(reader, ch) - return reader.schema, ch, nil + return doGetQuery(ctx, s.Alloc, s.db, string(cmd.GetStatementHandle()), nil) } func (s *SQLiteFlightSQLServer) GetFlightInfoCatalogs(_ context.Context, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { @@ -272,20 +297,7 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoTableTypes(_ context.Context, desc func (s *SQLiteFlightSQLServer) DoGetTableTypes(ctx context.Context) (*arrow.Schema, <-chan flight.StreamChunk, error) { query := "SELECT DISTINCT type AS table_type FROM sqlite_master" - - rows, err := s.db.QueryContext(ctx, query) - if err != nil { - return nil, nil, err - } - - reader, err := NewSqlBatchReaderWithSchema(s.Alloc, schema_ref.TableTypes, rows) - if err != nil { - return nil, nil, err - } - - ch := make(chan flight.StreamChunk) - go flight.StreamChunksFromReader(reader, ch) - return reader.schema, ch, nil + return doGetQuery(ctx, s.Alloc, s.db, query, schema_ref.TableTypes) } func (s *SQLiteFlightSQLServer) DoPutCommandStatementUpdate(ctx context.Context, cmd flightsql.StatementUpdate) (int64, error) { @@ -334,6 +346,28 @@ func (s *SQLiteFlightSQLServer) GetFlightInfoPreparedStatement(_ context.Context }, nil } +func doGetQuery(ctx context.Context, mem memory.Allocator, db *sql.DB, query string, schema *arrow.Schema, args ...interface{}) (*arrow.Schema, <-chan flight.StreamChunk, error) { + rows, err := db.QueryContext(ctx, query, args...) + if err != nil { + return nil, nil, err + } + + var rdr *SqlBatchReader + if schema != nil { + rdr, err = NewSqlBatchReaderWithSchema(mem, schema, rows) + } else { + rdr, err = NewSqlBatchReader(mem, rows) + } + + if err != nil { + return nil, nil, err + } + + ch := make(chan flight.StreamChunk) + go flight.StreamChunksFromReader(rdr, ch) + return rdr.schema, ch, nil +} + func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd flightsql.PreparedStatementQuery) (*arrow.Schema, <-chan flight.StreamChunk, error) { val, ok := s.prepared.Load(string(cmd.GetPreparedStatementHandle())) if !ok { @@ -436,3 +470,88 @@ func (s *SQLiteFlightSQLServer) DoPutPreparedStatementUpdate(ctx context.Context return result.RowsAffected() } + +func (s *SQLiteFlightSQLServer) GetFlightInfoPrimaryKeys(_ context.Context, cmd flightsql.TableRef, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.PrimaryKeys), nil +} + +func (s *SQLiteFlightSQLServer) DoGetPrimaryKeys(ctx context.Context, cmd flightsql.TableRef) (*arrow.Schema, <-chan flight.StreamChunk, error) { + // the field key_name can not be recovered by sqlite so it is + // being set to null following the same pattern for catalog name and schema_name + var b strings.Builder + + b.WriteString(` + SELECT null AS catalog_name, null AS schema_name, table_name, name AS column_name, pk AS key_sequence, null as key_name + FROM pragma_table_info(table_name) + JOIN (SELECT null AS catalog_name, null AS schema_name, name AS table_name, type AS table_type + FROM sqlite_master) where 1=1 AND pk !=0`) + + if cmd.Catalog != nil { + fmt.Fprintf(&b, " and catalog_name LIKE '%s'", *cmd.Catalog) + } + if cmd.DBSchema != nil { + fmt.Fprintf(&b, " and schema_name LIKE '%s'", *cmd.DBSchema) + } + + fmt.Fprintf(&b, " and table_name LIKE '%s'", cmd.Table) + + return doGetQuery(ctx, s.Alloc, s.db, b.String(), schema_ref.PrimaryKeys) +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoImportedKeys(_ context.Context, _ flightsql.TableRef, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.ImportedKeys), nil +} + +func (s *SQLiteFlightSQLServer) DoGetImportedKeys(ctx context.Context, ref flightsql.TableRef) (*arrow.Schema, <-chan flight.StreamChunk, error) { + filter := "fk_table_name = '" + ref.Table + "'" + if ref.Catalog != nil { + filter += " AND fk_catalog_name = '" + *ref.Catalog + "'" + } + if ref.DBSchema != nil { + filter += " AND fk_schema_name = '" + *ref.DBSchema + "'" + } + query := prepareQueryForGetKeys(filter) + return doGetQuery(ctx, s.Alloc, s.db, query, schema_ref.ImportedKeys) +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoExportedKeys(_ context.Context, _ flightsql.TableRef, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.ExportedKeys), nil +} + +func (s *SQLiteFlightSQLServer) DoGetExportedKeys(ctx context.Context, ref flightsql.TableRef) (*arrow.Schema, <-chan flight.StreamChunk, error) { + filter := "pk_table_name = '" + ref.Table + "'" + if ref.Catalog != nil { + filter += " AND pk_catalog_name = '" + *ref.Catalog + "'" + } + if ref.DBSchema != nil { + filter += " AND pk_schema_name = '" + *ref.DBSchema + "'" + } + query := prepareQueryForGetKeys(filter) + return doGetQuery(ctx, s.Alloc, s.db, query, schema_ref.ExportedKeys) +} + +func (s *SQLiteFlightSQLServer) GetFlightInfoCrossReference(_ context.Context, _ flightsql.CrossTableRef, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { + return s.flightInfoForCommand(desc, schema_ref.CrossReference), nil +} + +func (s *SQLiteFlightSQLServer) DoGetCrossReference(ctx context.Context, cmd flightsql.CrossTableRef) (*arrow.Schema, <-chan flight.StreamChunk, error) { + pkref := cmd.PKRef + filter := "pk_table_name = '" + pkref.Table + "'" + if pkref.Catalog != nil { + filter += " AND pk_catalog_name = '" + *pkref.Catalog + "'" + } + if pkref.DBSchema != nil { + filter += " AND pk_schema_name = '" + *pkref.DBSchema + "'" + } + + fkref := cmd.FKRef + filter += " AND fk_table_name = '" + fkref.Table + "'" + if fkref.Catalog != nil { + filter += " AND fk_catalog_name = '" + *fkref.Catalog + "'" + } + if fkref.DBSchema != nil { + filter += " AND fk_schema_name = '" + *fkref.DBSchema + "'" + } + query := prepareQueryForGetKeys(filter) + return doGetQuery(ctx, s.Alloc, s.db, query, schema_ref.ExportedKeys) +} diff --git a/go/arrow/flight/flightsql/server.go b/go/arrow/flight/flightsql/server.go index 3823b84c1d2ad..17bc9e188aa9c 100644 --- a/go/arrow/flight/flightsql/server.go +++ b/go/arrow/flight/flightsql/server.go @@ -164,6 +164,10 @@ func (BaseServer) mustEmbedBaseServer() {} // // Once registered, this value will be returned for any SqlInfo requests. func (b *BaseServer) RegisterSqlInfo(id SqlInfo, result interface{}) error { + if b.sqlInfoToResult == nil { + b.sqlInfoToResult = make(SqlInfoResultMap) + } + switch result.(type) { case string, bool, int64, int32, []string, map[int32][]int32: b.sqlInfoToResult[uint32(id)] = result diff --git a/go/arrow/flight/flightsql/server_test.go b/go/arrow/flight/flightsql/server_test.go index e488c802be6dd..9cfcb099f320c 100644 --- a/go/arrow/flight/flightsql/server_test.go +++ b/go/arrow/flight/flightsql/server_test.go @@ -31,6 +31,7 @@ import ( pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" "github.com/apache/arrow/go/v10/arrow/memory" "github.com/apache/arrow/go/v10/arrow/scalar" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -750,6 +751,213 @@ func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdate() { s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) } +func (s *FlightSqliteServerSuite) TestCommandGetPrimaryKeys() { + ctx := context.Background() + info, err := s.cl.GetPrimaryKeys(ctx, flightsql.TableRef{Table: "int%"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.PrimaryKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("intTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).(*array.Int32Builder).Append(1) + bldr.Field(5).AppendNull() + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetImportedKeys() { + ctx := context.Background() + info, err := s.cl.GetImportedKeys(ctx, flightsql.TableRef{Table: "intTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetExportedKeys() { + ctx := context.Background() + info, err := s.cl.GetExportedKeys(ctx, flightsql.TableRef{Table: "foreignTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetCrossRef() { + ctx := context.Background() + info, err := s.cl.GetCrossReference(ctx, + flightsql.TableRef{Table: "foreignTable"}, + flightsql.TableRef{Table: "intTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func validateSqlInfo(t *testing.T, expected interface{}, sc scalar.Scalar) bool { + switch ex := expected.(type) { + case string: + return assert.Equal(t, ex, sc.String()) + case bool: + return assert.Equal(t, ex, sc.(*scalar.Boolean).Value) + case int64: + return assert.Equal(t, ex, sc.(*scalar.Int64).Value) + case int32: + return assert.Equal(t, ex, sc.(*scalar.Int32).Value) + case []string: + arr := sc.(*scalar.List).Value.(*array.String) + assert.EqualValues(t, len(ex), arr.Len()) + for i, v := range ex { + assert.Equal(t, v, arr.Value(i)) + } + case map[int32][]int32: + // map is a list of structs with key and values + structArr := sc.(*scalar.Map).Value.(*array.Struct) + keys := structArr.Field(0).(*array.Int32) + values := structArr.Field(1).(*array.List) + // assert that the map has the right size + assert.EqualValues(t, len(ex), keys.Len()) + + // for each element, match the argument + for i := 0; i < keys.Len(); i++ { + keyScalar, _ := scalar.GetScalar(keys, i) + infoID := keyScalar.(*scalar.Int32).Value + + // assert the key exists + list, ok := ex[infoID] + assert.True(t, ok) + + // assert the int32list is the right size + start, end := values.ValueOffsets(i) + assert.EqualValues(t, len(list), end-start) + + // for each element make sure it matches + for j, v := range list { + listItem, err := scalar.GetScalar(values.ListValues(), int(start)+j) + assert.NoError(t, err) + assert.Equal(t, v, listItem.(*scalar.Int32).Value) + } + } + } + return true +} + +func (s *FlightSqliteServerSuite) TestCommandGetSqlInfo() { + expectedResults := example.SqlInfoResultMap() + infoIDs := make([]flightsql.SqlInfo, 0, len(expectedResults)) + for k := range expectedResults { + infoIDs = append(infoIDs, flightsql.SqlInfo(k)) + } + + ctx := context.Background() + info, err := s.cl.GetSqlInfo(ctx, infoIDs) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.EqualValues(2, rec.NumCols()) + s.EqualValues(len(expectedResults), rec.NumRows()) + + colName := rec.Column(0).(*array.Uint32) + colValue := rec.Column(1) + for i := 0; i < int(rec.NumRows()); i++ { + expected := expectedResults[colName.Value(i)] + sc, err := scalar.GetScalar(colValue, i) + s.NoError(err) + + s.True(validateSqlInfo(s.T(), expected, sc.(*scalar.DenseUnion).ChildValue())) + + sc.(*scalar.DenseUnion).Release() + } +} + func TestSqliteServer(t *testing.T) { suite.Run(t, new(FlightSqliteServerSuite)) } From 2dcf2d586fd1882a909c2535b27233cedd030479 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Sun, 14 Aug 2022 12:59:32 -0400 Subject: [PATCH 13/23] fix tests --- go/arrow/array/list.go | 5 +- .../flightsql/example/sql_batch_reader.go | 11 +- .../flight/flightsql/example/sqlite_test.go | 120 ------------------ go/go.mod | 2 +- go/go.sum | 8 +- 5 files changed, 17 insertions(+), 129 deletions(-) delete mode 100644 go/arrow/flight/flightsql/example/sqlite_test.go diff --git a/go/arrow/array/list.go b/go/arrow/array/list.go index 5b2e4b3f61f38..07e38944348ac 100644 --- a/go/arrow/array/list.go +++ b/go/arrow/array/list.go @@ -367,9 +367,9 @@ func NewLargeListBuilder(mem memory.Allocator, etype arrow.DataType) *LargeListB // NewLargeListBuilderWithField takes a field rather than just an element type // to allow for more customization of the final type of the LargeList Array -func NewLargeListBuilderWithField(mem memory.Allocator, field arrow.Field) *ListBuilder { +func NewLargeListBuilderWithField(mem memory.Allocator, field arrow.Field) *LargeListBuilder { offsetBldr := NewInt64Builder(mem) - return &ListBuilder{ + return &LargeListBuilder{ baseListBuilder{ builder: builder{refCount: 1, mem: mem}, values: NewBuilder(mem, field.Type), @@ -398,7 +398,6 @@ func (b *baseListBuilder) Release() { func (b *baseListBuilder) appendNextOffset() { b.appendOffsetVal(b.values.Len()) - // b.offsets.Append(int32(b.values.Len())) } func (b *baseListBuilder) Append(v bool) { diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index d56f057a477c3..4456b5c38d96d 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -82,7 +82,7 @@ func NewSqlBatchReaderWithSchema(mem memory.Allocator, schema *arrow.Schema, row switch f.Type.ID() { case arrow.UINT8: if f.Nullable { - rowdest[i] = &sql.NullInt16{} + rowdest[i] = &sql.NullInt32{} } else { rowdest[i] = new(uint8) } @@ -143,7 +143,7 @@ func NewSqlBatchReader(mem memory.Allocator, rows *sql.Rows) (*SqlBatchReader, e switch fields[i].Type.ID() { case arrow.UINT8: if fields[i].Nullable { - rowdest[i] = &sql.NullInt16{} + rowdest[i] = &sql.NullInt32{} } else { rowdest[i] = new(uint8) } @@ -248,7 +248,12 @@ func (r *SqlBatchReader) Next() bool { if !v.Valid { fb.AppendNull() } else { - fb.(*array.Int32Builder).Append(v.Int32) + switch b := fb.(type) { + case *array.Int32Builder: + b.Append(v.Int32) + case *array.Uint8Builder: + b.Append(uint8(v.Int32)) + } } case *float64: fb.(*array.Float64Builder).Append(*v) diff --git a/go/arrow/flight/flightsql/example/sqlite_test.go b/go/arrow/flight/flightsql/example/sqlite_test.go deleted file mode 100644 index b8e9be831ffbd..0000000000000 --- a/go/arrow/flight/flightsql/example/sqlite_test.go +++ /dev/null @@ -1,120 +0,0 @@ -// 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. - -package example - -import ( - "context" - "strings" - "testing" - - "github.com/apache/arrow/go/v10/arrow" - "github.com/apache/arrow/go/v10/arrow/array" - "github.com/apache/arrow/go/v10/arrow/flight" - "github.com/apache/arrow/go/v10/arrow/flight/flightsql" - "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" - "github.com/apache/arrow/go/v10/arrow/memory" - "github.com/apache/arrow/go/v10/arrow/scalar" - "github.com/stretchr/testify/assert" - "google.golang.org/grpc" - "google.golang.org/grpc/credentials/insecure" -) - -var dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} - -func TestServer(t *testing.T) { - mem := memory.NewCheckedAllocator(memory.DefaultAllocator) - defer mem.AssertSize(t, 0) - - s := flight.NewServerWithMiddleware(nil) - srv, err := NewSQLiteFlightSQLServer() - assert.NoError(t, err) - assert.NotNil(t, srv) - srv.Alloc = mem - s.RegisterFlightService(flightsql.NewFlightServer(srv)) - s.Init("localhost:0") - - go s.Serve() - defer s.Shutdown() - - cl, err := flightsql.NewClient(s.Addr().String(), nil, nil, dialOpts...) - assert.NoError(t, err) - assert.NotNil(t, cl) - - info, err := cl.Execute(context.Background(), "SELECT * FROM intTable") - assert.NoError(t, err) - - rdr, err := cl.DoGet(context.Background(), info.Endpoint[0].Ticket) - assert.NoError(t, err) - defer rdr.Release() - - assert.True(t, rdr.Next()) - rec := rdr.Record() - assert.NotNil(t, rec) - - expectedSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "keyName", Type: arrow.BinaryTypes.String, Nullable: true}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Nullable: true}, - }, nil) - - assert.Truef(t, expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) - - idarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 2, 3, 4]`)) - defer idarr.Release() - keyarr, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["one", "zero", "negative one", null]`)) - defer keyarr.Release() - valarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 0, -1, null]`)) - defer valarr.Release() - foreignarr, _, _ := array.FromJSON(mem, arrow.PrimitiveTypes.Int64, strings.NewReader(`[1, 1, 1, null]`)) - defer foreignarr.Release() - - expectedRec := array.NewRecord(expectedSchema, []arrow.Array{idarr, keyarr, valarr, foreignarr}, 4) - defer expectedRec.Release() - - assert.Truef(t, array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) - - info, err = cl.GetTables(context.Background(), &flightsql.GetTablesOpts{}) - assert.NoError(t, err) - assert.NotNil(t, info) - - rdr, err = cl.DoGet(context.Background(), info.Endpoint[0].Ticket) - assert.NoError(t, err) - defer rdr.Release() - - catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, mem) - defer catalogName.Release() - schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, mem) - defer schemaName.Release() - - tableName, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["foreignTable", "intTable", "sqlite_sequence"]`)) - defer tableName.Release() - - tableType, _, _ := array.FromJSON(mem, arrow.BinaryTypes.String, strings.NewReader(`["table", "table", "table"]`)) - defer tableType.Release() - - expectedRec = array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) - defer expectedRec.Release() - - assert.True(t, rdr.Next()) - rec = rdr.Record() - assert.NotNil(t, rec) - rec.Retain() - assert.False(t, rdr.Next()) - - assert.Truef(t, array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) -} diff --git a/go/go.mod b/go/go.mod index 6a41ad0f06a17..6abb73a45d297 100644 --- a/go/go.mod +++ b/go/go.mod @@ -50,6 +50,6 @@ require ( modernc.org/ccgo/v3 v3.16.8 // indirect modernc.org/libc v1.16.19 // indirect modernc.org/opt v0.1.3 // indirect - modernc.org/sqlite v1.18.0 // indirect + modernc.org/sqlite v1.18.0 modernc.org/strutil v1.1.2 // indirect ) diff --git a/go/go.sum b/go/go.sum index c2973d283fa67..58e2996e4351f 100644 --- a/go/go.sum +++ b/go/go.sum @@ -32,6 +32,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815 h1:bWDMxwH3px2JBh6AyO7hdCn/PkvCZXii8TGj7sbtEbQ= github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE= +github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= @@ -110,6 +111,7 @@ github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= +github.com/mattn/go-sqlite3 v1.14.12 h1:TJ1bhYJPV44phC+IMu1u2K/i5RriLTPe+yc68XDJ1Z0= github.com/mattn/go-sqlite3 v1.14.12/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= @@ -223,8 +225,6 @@ golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20211019181941-9d821ace8654/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20220804214406-8e32c043e418 h1:9vYwv7OjYaky/tlAeD7C4oC9EsPTlaFl1H2jS++V+ME= -golang.org/x/sys v0.0.0-20220804214406-8e32c043e418/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220808155132-1c4a2a72c664 h1:v1W7bwXHsnLLloWYTVEdvGvA7BHMeBYsPcF0GLDxIRs= golang.org/x/sys v0.0.0-20220808155132-1c4a2a72c664/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= @@ -317,7 +317,9 @@ modernc.org/ccgo/v3 v3.16.4/go.mod h1:tGtX0gE9Jn7hdZFeU88slbTh1UtCYKusWOoCJuvkWs modernc.org/ccgo/v3 v3.16.6/go.mod h1:tGtX0gE9Jn7hdZFeU88slbTh1UtCYKusWOoCJuvkWsQ= modernc.org/ccgo/v3 v3.16.8 h1:G0QNlTqI5uVgczBWfGKs7B++EPwCfXPWGD2MdeKloDs= modernc.org/ccgo/v3 v3.16.8/go.mod h1:zNjwkizS+fIFDrDjIAgBSCLkWbJuHF+ar3QRn+Z9aws= +modernc.org/ccorpus v1.11.6 h1:J16RXiiqiCgua6+ZvQot4yUuUy8zxgqbqEEUuGPlISk= modernc.org/ccorpus v1.11.6/go.mod h1:2gEUTrWqdpH2pXsmTM1ZkjeSrUWDpjMu2T6m29L/ErQ= +modernc.org/httpfs v1.0.6 h1:AAgIpFZRXuYnkjftxTAZwMIiwEqAfk8aVB2/oA6nAeM= modernc.org/httpfs v1.0.6/go.mod h1:7dosgurJGp0sPaRanU53W4xZYKh14wfzX420oZADeHM= modernc.org/libc v0.0.0-20220428101251-2d5f3daf273b/go.mod h1:p7Mg4+koNjc8jkqwcoFBJx7tXkpj00G77X7A72jXPXA= modernc.org/libc v1.16.0/go.mod h1:N4LD6DBE9cf+Dzf9buBlzVJndKr/iJHG97vGLHYnb5A= @@ -339,8 +341,10 @@ modernc.org/sqlite v1.18.0/go.mod h1:B9fRWZacNxJBHoCJZQr1R54zhVn3fjfl0aszflrTSxY modernc.org/strutil v1.1.1/go.mod h1:DE+MQQ/hjKBZS2zNInV5hhcipt5rLPWkmpbGeW5mmdw= modernc.org/strutil v1.1.2 h1:iFBDH6j1Z0bN/Q9udJnnFoFpENA4252qe/7/5woE5MI= modernc.org/strutil v1.1.2/go.mod h1:OYajnUAcI/MX+XD/Wx7v1bbdvcQSvxgtb0gC+u3d3eg= +modernc.org/tcl v1.13.1 h1:npxzTwFTZYM8ghWicVIX1cRWzj7Nd8i6AqqX2p+IYao= modernc.org/tcl v1.13.1/go.mod h1:XOLfOwzhkljL4itZkK6T72ckMgvj0BDsnKNdZVUOecw= modernc.org/token v1.0.0 h1:a0jaWiNMDhDUtqOj09wvjWWAqd3q7WpBulmL9H2egsk= modernc.org/token v1.0.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= +modernc.org/z v1.5.1 h1:RTNHdsrOpeoSeOF4FbzTo8gBYByaJ5xT7NgZ9ZqRiJM= modernc.org/z v1.5.1/go.mod h1:eWFB510QWW5Th9YGZT81s+LwvaAs3Q2yr4sP0rmLkv8= rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= From 84d5d073c039601b10463d72806f6b44e34fd79c Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Sun, 14 Aug 2022 13:03:50 -0400 Subject: [PATCH 14/23] nullint16 isn't in go1.16 --- go/arrow/flight/flightsql/example/sql_batch_reader.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index 4456b5c38d96d..998da90234187 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -228,12 +228,6 @@ func (r *SqlBatchReader) Next() bool { switch v := v.(type) { case *uint8: fb.(*array.Uint8Builder).Append(*v) - case *sql.NullInt16: - if !v.Valid { - fb.AppendNull() - } else { - fb.(*array.Uint8Builder).Append(uint8(v.Int16)) - } case *int64: fb.(*array.Int64Builder).Append(*v) case *sql.NullInt64: From 093b7ef5c973ec77c852c498c6ddc11b55cc2c88 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Sun, 14 Aug 2022 13:17:36 -0400 Subject: [PATCH 15/23] alleviate race conditions --- go/arrow/flight/flightsql/example/sqlite_server.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go index 40e1c0979988c..c4b83b1af5136 100644 --- a/go/arrow/flight/flightsql/example/sqlite_server.go +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -269,8 +269,9 @@ func (s *SQLiteFlightSQLServer) DoGetTables(ctx context.Context, cmd flightsql.G } } + schema := rdr.Schema() go flight.StreamChunksFromReader(rdr, ch) - return rdr.Schema(), ch, nil + return schema, ch, nil } func (s *SQLiteFlightSQLServer) GetFlightInfoXdbcTypeInfo(_ context.Context, _ flightsql.GetXdbcTypeInfo, desc *flight.FlightDescriptor) (*flight.FlightInfo, error) { @@ -357,6 +358,9 @@ func doGetQuery(ctx context.Context, mem memory.Allocator, db *sql.DB, query str rdr, err = NewSqlBatchReaderWithSchema(mem, schema, rows) } else { rdr, err = NewSqlBatchReader(mem, rows) + if err == nil { + schema = rdr.schema + } } if err != nil { @@ -365,7 +369,7 @@ func doGetQuery(ctx context.Context, mem memory.Allocator, db *sql.DB, query str ch := make(chan flight.StreamChunk) go flight.StreamChunksFromReader(rdr, ch) - return rdr.schema, ch, nil + return schema, ch, nil } func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd flightsql.PreparedStatementQuery) (*arrow.Schema, <-chan flight.StreamChunk, error) { @@ -385,9 +389,10 @@ func (s *SQLiteFlightSQLServer) DoGetPreparedStatement(ctx context.Context, cmd return nil, nil, err } + schema := rdr.schema ch := make(chan flight.StreamChunk) go flight.StreamChunksFromReader(rdr, ch) - return rdr.Schema(), ch, nil + return schema, ch, nil } func getParamsForStatement(rdr flight.MessageReader) (params []interface{}, err error) { From 36bbd170252a961bd9ed762709b47daf04eaadfc Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Sun, 14 Aug 2022 13:53:10 -0400 Subject: [PATCH 16/23] put sqlite example server behind go1.17+ --- .../flightsql/example/sql_batch_reader.go | 3 + .../flight/flightsql/example/sqlite_info.go | 3 + .../flight/flightsql/example/sqlite_server.go | 7 + .../sqlite_tables_schema_batch_reader.go | 3 + .../flight/flightsql/example/type_info.go | 3 + go/arrow/flight/flightsql/server_test.go | 751 ----------------- .../flight/flightsql/sqlite_server_test.go | 783 ++++++++++++++++++ 7 files changed, 802 insertions(+), 751 deletions(-) create mode 100644 go/arrow/flight/flightsql/sqlite_server_test.go diff --git a/go/arrow/flight/flightsql/example/sql_batch_reader.go b/go/arrow/flight/flightsql/example/sql_batch_reader.go index 998da90234187..5e7fbde1afef7 100644 --- a/go/arrow/flight/flightsql/example/sql_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sql_batch_reader.go @@ -14,6 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +//go:build go1.17 +// +build go1.17 + package example import ( diff --git a/go/arrow/flight/flightsql/example/sqlite_info.go b/go/arrow/flight/flightsql/example/sqlite_info.go index bc7f34c56e571..5a85b21412866 100644 --- a/go/arrow/flight/flightsql/example/sqlite_info.go +++ b/go/arrow/flight/flightsql/example/sqlite_info.go @@ -14,6 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +//go:build go1.17 +// +build go1.17 + package example import "github.com/apache/arrow/go/v10/arrow/flight/flightsql" diff --git a/go/arrow/flight/flightsql/example/sqlite_server.go b/go/arrow/flight/flightsql/example/sqlite_server.go index c4b83b1af5136..5d2599b52c09f 100644 --- a/go/arrow/flight/flightsql/example/sqlite_server.go +++ b/go/arrow/flight/flightsql/example/sqlite_server.go @@ -14,6 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +//go:build go1.17 +// +build go1.17 + // Package example contains a FlightSQL Server implementation using // sqlite as the backing engine. // @@ -26,6 +29,10 @@ // That said, since both implement in terms of Go's standard database/sql // package, it's easy to swap them out if desired as the modernc.org/sqlite // package is slower than go-sqlite3. +// +// One other important note is that modernc.org/sqlite only works in go +// 1.17+ so this entire package is given the build constraint to only +// build when using go1.17 or higher package example import ( diff --git a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go index 005b7e5317e0b..851b301c7482f 100644 --- a/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go +++ b/go/arrow/flight/flightsql/example/sqlite_tables_schema_batch_reader.go @@ -14,6 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +//go:build go1.17 +// +build go1.17 + package example import ( diff --git a/go/arrow/flight/flightsql/example/type_info.go b/go/arrow/flight/flightsql/example/type_info.go index 9e301bddf9e57..dcba42b1f847b 100644 --- a/go/arrow/flight/flightsql/example/type_info.go +++ b/go/arrow/flight/flightsql/example/type_info.go @@ -14,6 +14,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +//go:build go1.17 +// +build go1.17 + package example import ( diff --git a/go/arrow/flight/flightsql/server_test.go b/go/arrow/flight/flightsql/server_test.go index 9cfcb099f320c..ece7754bbbdb2 100644 --- a/go/arrow/flight/flightsql/server_test.go +++ b/go/arrow/flight/flightsql/server_test.go @@ -18,20 +18,13 @@ package flightsql_test import ( "context" - "os" "strings" "testing" - "github.com/apache/arrow/go/v10/arrow" - "github.com/apache/arrow/go/v10/arrow/array" "github.com/apache/arrow/go/v10/arrow/flight" "github.com/apache/arrow/go/v10/arrow/flight/flightsql" - "github.com/apache/arrow/go/v10/arrow/flight/flightsql/example" - "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" pb "github.com/apache/arrow/go/v10/arrow/flight/internal/flight" "github.com/apache/arrow/go/v10/arrow/memory" - "github.com/apache/arrow/go/v10/arrow/scalar" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -39,7 +32,6 @@ import ( "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/anypb" - sqlite3 "modernc.org/sqlite/lib" ) var dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} @@ -218,746 +210,3 @@ func (s *UnimplementedFlightSqlServerSuite) TestDoAction() { func TestBaseServer(t *testing.T) { suite.Run(t, new(UnimplementedFlightSqlServerSuite)) } - -type FlightSqliteServerSuite struct { - suite.Suite - - srv *example.SQLiteFlightSQLServer - s flight.Server - cl *flightsql.Client - - mem *memory.CheckedAllocator -} - -func (s *FlightSqliteServerSuite) getColMetadata(colType int, table string) arrow.Metadata { - bldr := flightsql.NewColumnMetadataBuilder() - bldr.Scale(15).IsReadOnly(false).IsAutoIncrement(false) - if table != "" { - bldr.TableName(table) - } - switch colType { - case sqlite3.SQLITE_TEXT, sqlite3.SQLITE_BLOB: - case sqlite3.SQLITE_INTEGER: - bldr.Precision(10) - case sqlite3.SQLITE_FLOAT: - bldr.Precision(15) - default: - bldr.Precision(0) - } - return bldr.Metadata() -} - -func (s *FlightSqliteServerSuite) SetupTest() { - var err error - s.mem = memory.NewCheckedAllocator(memory.DefaultAllocator) - s.s = flight.NewServerWithMiddleware(nil) - s.srv, err = example.NewSQLiteFlightSQLServer() - s.Require().NoError(err) - s.srv.Alloc = s.mem - - s.s.RegisterFlightService(flightsql.NewFlightServer(s.srv)) - s.s.Init("localhost:0") - s.s.SetShutdownOnSignals(os.Interrupt, os.Kill) - go s.s.Serve() - s.cl, err = flightsql.NewClient(s.s.Addr().String(), nil, nil, dialOpts...) - s.Require().NoError(err) - s.Require().NotNil(s.cl) - s.cl.Alloc = s.mem -} - -func (s *FlightSqliteServerSuite) TearDownTest() { - s.Require().NoError(s.cl.Close()) - s.s.Shutdown() - s.srv = nil - s.mem.AssertSize(s.T(), 0) -} - -func (s *FlightSqliteServerSuite) fromJSON(dt arrow.DataType, json string) arrow.Array { - arr, _, _ := array.FromJSON(s.mem, dt, strings.NewReader(json)) - return arr -} - -func (s *FlightSqliteServerSuite) execCountQuery(query string) int64 { - info, err := s.cl.Execute(context.Background(), query) - s.NoError(err) - - rdr, err := s.cl.DoGet(context.Background(), info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - rec, err := rdr.Read() - s.NoError(err) - return rec.Column(0).(*array.Int64).Value(0) -} - -func (s *FlightSqliteServerSuite) TestCommandStatementQuery() { - ctx := context.Background() - info, err := s.cl.Execute(ctx, "SELECT * FROM intTable") - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.NotNil(rec) - - expectedSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - }, nil) - - s.Truef(expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) - - idarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) - defer idarr.Release() - keyarr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) - defer keyarr.Release() - valarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) - defer valarr.Release() - foreignarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) - defer foreignarr.Release() - - expectedRec := array.NewRecord(expectedSchema, []arrow.Array{idarr, keyarr, valarr, foreignarr}, 4) - defer expectedRec.Release() - - s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTables() { - ctx := context.Background() - info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{}) - s.NoError(err) - s.NotNil(info) - - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) - defer catalogName.Release() - schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) - defer schemaName.Release() - - tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) - defer tableName.Release() - - tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) - defer tableType.Release() - - expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) - defer expectedRec.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.NotNil(rec) - rec.Retain() - defer rec.Release() - s.False(rdr.Next()) - - s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableFilter() { - ctx := context.Background() - info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ - TableNameFilterPattern: proto.String("int%"), - }) - s.NoError(err) - s.NotNil(info) - - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) - schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) - table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) - tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) - expected := array.NewRecord(schema_ref.Tables, []arrow.Array{catalog, schema, table, tabletype}, 1) - defer func() { - catalog.Release() - schema.Release() - table.Release() - tabletype.Release() - expected.Release() - }() - - s.True(rdr.Next()) - rec := rdr.Record() - s.NotNil(rec) - rec.Retain() - defer rec.Release() - s.False(rdr.Next()) - s.NoError(rdr.Err()) - - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableTypesFilter() { - ctx := context.Background() - info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ - TableTypes: []string{"index"}, - }) - s.NoError(err) - - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - s.True(schema_ref.Tables.Equal(rdr.Schema()), rdr.Schema().String()) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTablesWithExistingTableTypeFilter() { - ctx := context.Background() - info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ - TableTypes: []string{"table"}, - }) - s.NoError(err) - s.NotNil(info) - - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) - defer catalogName.Release() - schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) - defer schemaName.Release() - - tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) - defer tableName.Release() - - tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) - defer tableType.Release() - - expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) - defer expectedRec.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.NotNil(rec) - rec.Retain() - defer rec.Release() - s.False(rdr.Next()) - - s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTablesWithIncludedSchemas() { - ctx := context.Background() - info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ - TableNameFilterPattern: proto.String("int%"), - IncludeSchema: true, - }) - s.NoError(err) - s.NotNil(info) - - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) - schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) - table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) - tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) - - dbTableName := "intTable" - - tableSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, - Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, - {Name: "keyName", Type: arrow.BinaryTypes.String, - Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, dbTableName)}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, - Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, - Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, - }, nil) - schemaBuf := flight.SerializeSchema(tableSchema, s.mem) - binaryBldr := array.NewBinaryBuilder(s.mem, arrow.BinaryTypes.Binary) - binaryBldr.Append(schemaBuf) - schemaCol := binaryBldr.NewArray() - - expected := array.NewRecord(schema_ref.TablesWithIncludedSchema, []arrow.Array{catalog, schema, table, tabletype, schemaCol}, 1) - defer func() { - catalog.Release() - schema.Release() - table.Release() - tabletype.Release() - binaryBldr.Release() - schemaCol.Release() - expected.Release() - }() - - s.True(rdr.Next()) - rec := rdr.Record() - s.NotNil(rec) - rec.Retain() - defer rec.Release() - s.False(rdr.Next()) - s.NoError(rdr.Err()) - - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTypeInfo() { - ctx := context.Background() - info, err := s.cl.GetXdbcTypeInfo(ctx, nil) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - expected := example.GetTypeInfoResult(s.mem) - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTypeInfoFiltered() { - ctx := context.Background() - info, err := s.cl.GetXdbcTypeInfo(ctx, proto.Int32(-4)) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - expected := example.GetFilteredTypeInfoResult(s.mem, -4) - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetCatalogs() { - ctx := context.Background() - info, err := s.cl.GetCatalogs(ctx) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - s.True(rdr.Schema().Equal(schema_ref.Catalogs), rdr.Schema().String()) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetDbSchemas() { - ctx := context.Background() - info, err := s.cl.GetDBSchemas(ctx, &flightsql.GetDBSchemasOpts{}) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - s.True(rdr.Schema().Equal(schema_ref.DBSchemas), rdr.Schema().String()) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetTableTypes() { - ctx := context.Background() - info, err := s.cl.GetTableTypes(ctx) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - expected := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) - defer expected.Release() - expectedRec := array.NewRecord(schema_ref.TableTypes, []arrow.Array{expected}, 1) - defer expectedRec.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandStatementUpdate() { - ctx := context.Background() - result, err := s.cl.ExecuteUpdate(ctx, `INSERT INTO intTable (keyName, value) VALUES - ('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)`) - s.NoError(err) - s.EqualValues(3, result) - - result, err = s.cl.ExecuteUpdate(ctx, `UPDATE intTable SET keyName = 'KEYNAME1' - WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'`) - s.NoError(err) - s.EqualValues(2, result) - - result, err = s.cl.ExecuteUpdate(ctx, `DELETE FROM intTable WHERE keyName = 'KEYNAME1'`) - s.NoError(err) - s.EqualValues(3, result) -} - -func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQuery() { - ctx := context.Background() - prep, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable") - s.NoError(err) - defer prep.Close(ctx) - - info, err := prep.Execute(ctx) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - - expectedSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) - - idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) - defer idArr.Release() - keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) - defer keyNameArr.Release() - valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) - defer valueArr.Release() - foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) - defer foreignIdArr.Release() - - expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 4) - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQueryWithParams() { - ctx := context.Background() - stmt, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable WHERE keyName LIKE ?") - s.NoError(err) - defer stmt.Close(ctx) - - typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[0]") - offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") - strArray := s.fromJSON(arrow.BinaryTypes.String, `["%one"]`) - bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") - bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[]") - dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") - paramArr, _ := array.NewDenseUnionFromArraysWithFields(typeIDs, - offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, - []string{"string", "bytes", "bigint", "double"}) - batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ - {Name: "parameter_1", Type: paramArr.DataType()}}, nil), - []arrow.Array{paramArr}, 1) - defer func() { - typeIDs.Release() - offsets.Release() - strArray.Release() - bytesArr.Release() - bigintArr.Release() - dblArr.Release() - paramArr.Release() - batch.Release() - }() - - stmt.SetParameters(batch) - info, err := stmt.Execute(ctx) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - - expectedSchema := arrow.NewSchema([]arrow.Field{ - {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, - {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, - {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) - - idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 3]`) - defer idArr.Release() - keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "negative one"]`) - defer keyNameArr.Release() - valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, -1]`) - defer valueArr.Release() - foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1]`) - defer foreignIdArr.Release() - - expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 2) - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdateWithParams() { - ctx := context.Background() - stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', ?)") - s.NoError(err) - defer stmt.Close(ctx) - - typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[2]") - offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") - strArray := s.fromJSON(arrow.BinaryTypes.String, "[]") - bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") - bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[999]") - dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") - paramArr, err := array.NewDenseUnionFromArraysWithFields(typeIDs, - offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, - []string{"string", "bytes", "bigint", "double"}) - s.NoError(err) - batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ - {Name: "parameter_1", Type: paramArr.DataType()}}, nil), - []arrow.Array{paramArr}, 1) - defer func() { - typeIDs.Release() - offsets.Release() - strArray.Release() - bytesArr.Release() - bigintArr.Release() - dblArr.Release() - paramArr.Release() - batch.Release() - }() - - stmt.SetParameters(batch) - s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) - n, err := stmt.ExecuteUpdate(context.Background()) - s.NoError(err) - s.EqualValues(1, n) - s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) - n, err = s.cl.ExecuteUpdate(context.Background(), "DELETE FROM intTable WHERE keyName = 'new_value'") - s.NoError(err) - s.EqualValues(1, n) - s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) -} - -func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdate() { - ctx := context.Background() - stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', 999)") - s.NoError(err) - defer stmt.Close(ctx) - - s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) - result, err := stmt.ExecuteUpdate(ctx) - s.NoError(err) - s.EqualValues(1, result) - s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) - result, err = s.cl.ExecuteUpdate(ctx, "DELETE FROM intTable WHERE keyName = 'new_value'") - s.NoError(err) - s.EqualValues(1, result) - s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) -} - -func (s *FlightSqliteServerSuite) TestCommandGetPrimaryKeys() { - ctx := context.Background() - info, err := s.cl.GetPrimaryKeys(ctx, flightsql.TableRef{Table: "int%"}) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - bldr := array.NewRecordBuilder(s.mem, schema_ref.PrimaryKeys) - defer bldr.Release() - bldr.Field(0).AppendNull() - bldr.Field(1).AppendNull() - bldr.Field(2).(*array.StringBuilder).Append("intTable") - bldr.Field(3).(*array.StringBuilder).Append("id") - bldr.Field(4).(*array.Int32Builder).Append(1) - bldr.Field(5).AppendNull() - expected := bldr.NewRecord() - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetImportedKeys() { - ctx := context.Background() - info, err := s.cl.GetImportedKeys(ctx, flightsql.TableRef{Table: "intTable"}) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) - defer bldr.Release() - bldr.Field(0).AppendNull() - bldr.Field(1).AppendNull() - bldr.Field(2).(*array.StringBuilder).Append("foreignTable") - bldr.Field(3).(*array.StringBuilder).Append("id") - bldr.Field(4).AppendNull() - bldr.Field(5).AppendNull() - bldr.Field(6).(*array.StringBuilder).Append("intTable") - bldr.Field(7).(*array.StringBuilder).Append("foreignId") - bldr.Field(8).(*array.Int32Builder).Append(0) - bldr.Field(9).AppendNull() - bldr.Field(10).AppendNull() - bldr.Field(11).(*array.Uint8Builder).Append(3) - bldr.Field(12).(*array.Uint8Builder).Append(3) - expected := bldr.NewRecord() - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetExportedKeys() { - ctx := context.Background() - info, err := s.cl.GetExportedKeys(ctx, flightsql.TableRef{Table: "foreignTable"}) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) - defer bldr.Release() - bldr.Field(0).AppendNull() - bldr.Field(1).AppendNull() - bldr.Field(2).(*array.StringBuilder).Append("foreignTable") - bldr.Field(3).(*array.StringBuilder).Append("id") - bldr.Field(4).AppendNull() - bldr.Field(5).AppendNull() - bldr.Field(6).(*array.StringBuilder).Append("intTable") - bldr.Field(7).(*array.StringBuilder).Append("foreignId") - bldr.Field(8).(*array.Int32Builder).Append(0) - bldr.Field(9).AppendNull() - bldr.Field(10).AppendNull() - bldr.Field(11).(*array.Uint8Builder).Append(3) - bldr.Field(12).(*array.Uint8Builder).Append(3) - expected := bldr.NewRecord() - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func (s *FlightSqliteServerSuite) TestCommandGetCrossRef() { - ctx := context.Background() - info, err := s.cl.GetCrossReference(ctx, - flightsql.TableRef{Table: "foreignTable"}, - flightsql.TableRef{Table: "intTable"}) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) - defer bldr.Release() - bldr.Field(0).AppendNull() - bldr.Field(1).AppendNull() - bldr.Field(2).(*array.StringBuilder).Append("foreignTable") - bldr.Field(3).(*array.StringBuilder).Append("id") - bldr.Field(4).AppendNull() - bldr.Field(5).AppendNull() - bldr.Field(6).(*array.StringBuilder).Append("intTable") - bldr.Field(7).(*array.StringBuilder).Append("foreignId") - bldr.Field(8).(*array.Int32Builder).Append(0) - bldr.Field(9).AppendNull() - bldr.Field(10).AppendNull() - bldr.Field(11).(*array.Uint8Builder).Append(3) - bldr.Field(12).(*array.Uint8Builder).Append(3) - expected := bldr.NewRecord() - defer expected.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) - s.False(rdr.Next()) -} - -func validateSqlInfo(t *testing.T, expected interface{}, sc scalar.Scalar) bool { - switch ex := expected.(type) { - case string: - return assert.Equal(t, ex, sc.String()) - case bool: - return assert.Equal(t, ex, sc.(*scalar.Boolean).Value) - case int64: - return assert.Equal(t, ex, sc.(*scalar.Int64).Value) - case int32: - return assert.Equal(t, ex, sc.(*scalar.Int32).Value) - case []string: - arr := sc.(*scalar.List).Value.(*array.String) - assert.EqualValues(t, len(ex), arr.Len()) - for i, v := range ex { - assert.Equal(t, v, arr.Value(i)) - } - case map[int32][]int32: - // map is a list of structs with key and values - structArr := sc.(*scalar.Map).Value.(*array.Struct) - keys := structArr.Field(0).(*array.Int32) - values := structArr.Field(1).(*array.List) - // assert that the map has the right size - assert.EqualValues(t, len(ex), keys.Len()) - - // for each element, match the argument - for i := 0; i < keys.Len(); i++ { - keyScalar, _ := scalar.GetScalar(keys, i) - infoID := keyScalar.(*scalar.Int32).Value - - // assert the key exists - list, ok := ex[infoID] - assert.True(t, ok) - - // assert the int32list is the right size - start, end := values.ValueOffsets(i) - assert.EqualValues(t, len(list), end-start) - - // for each element make sure it matches - for j, v := range list { - listItem, err := scalar.GetScalar(values.ListValues(), int(start)+j) - assert.NoError(t, err) - assert.Equal(t, v, listItem.(*scalar.Int32).Value) - } - } - } - return true -} - -func (s *FlightSqliteServerSuite) TestCommandGetSqlInfo() { - expectedResults := example.SqlInfoResultMap() - infoIDs := make([]flightsql.SqlInfo, 0, len(expectedResults)) - for k := range expectedResults { - infoIDs = append(infoIDs, flightsql.SqlInfo(k)) - } - - ctx := context.Background() - info, err := s.cl.GetSqlInfo(ctx, infoIDs) - s.NoError(err) - rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) - s.NoError(err) - defer rdr.Release() - - s.True(rdr.Next()) - rec := rdr.Record() - rec.Retain() - defer rec.Release() - s.False(rdr.Next()) - - s.EqualValues(2, rec.NumCols()) - s.EqualValues(len(expectedResults), rec.NumRows()) - - colName := rec.Column(0).(*array.Uint32) - colValue := rec.Column(1) - for i := 0; i < int(rec.NumRows()); i++ { - expected := expectedResults[colName.Value(i)] - sc, err := scalar.GetScalar(colValue, i) - s.NoError(err) - - s.True(validateSqlInfo(s.T(), expected, sc.(*scalar.DenseUnion).ChildValue())) - - sc.(*scalar.DenseUnion).Release() - } -} - -func TestSqliteServer(t *testing.T) { - suite.Run(t, new(FlightSqliteServerSuite)) -} diff --git a/go/arrow/flight/flightsql/sqlite_server_test.go b/go/arrow/flight/flightsql/sqlite_server_test.go new file mode 100644 index 0000000000000..b6e6335700d83 --- /dev/null +++ b/go/arrow/flight/flightsql/sqlite_server_test.go @@ -0,0 +1,783 @@ +// 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. + +//go:build go1.17 +// +build go1.17 + +package flightsql_test + +import ( + "context" + "os" + "strings" + "testing" + + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/array" + "github.com/apache/arrow/go/v10/arrow/flight" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/example" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql/schema_ref" + "github.com/apache/arrow/go/v10/arrow/memory" + "github.com/apache/arrow/go/v10/arrow/scalar" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/suite" + "google.golang.org/protobuf/proto" + sqlite3 "modernc.org/sqlite/lib" +) + +type FlightSqliteServerSuite struct { + suite.Suite + + srv *example.SQLiteFlightSQLServer + s flight.Server + cl *flightsql.Client + + mem *memory.CheckedAllocator +} + +func (s *FlightSqliteServerSuite) getColMetadata(colType int, table string) arrow.Metadata { + bldr := flightsql.NewColumnMetadataBuilder() + bldr.Scale(15).IsReadOnly(false).IsAutoIncrement(false) + if table != "" { + bldr.TableName(table) + } + switch colType { + case sqlite3.SQLITE_TEXT, sqlite3.SQLITE_BLOB: + case sqlite3.SQLITE_INTEGER: + bldr.Precision(10) + case sqlite3.SQLITE_FLOAT: + bldr.Precision(15) + default: + bldr.Precision(0) + } + return bldr.Metadata() +} + +func (s *FlightSqliteServerSuite) SetupTest() { + var err error + s.mem = memory.NewCheckedAllocator(memory.DefaultAllocator) + s.s = flight.NewServerWithMiddleware(nil) + s.srv, err = example.NewSQLiteFlightSQLServer() + s.Require().NoError(err) + s.srv.Alloc = s.mem + + s.s.RegisterFlightService(flightsql.NewFlightServer(s.srv)) + s.s.Init("localhost:0") + s.s.SetShutdownOnSignals(os.Interrupt, os.Kill) + go s.s.Serve() + s.cl, err = flightsql.NewClient(s.s.Addr().String(), nil, nil, dialOpts...) + s.Require().NoError(err) + s.Require().NotNil(s.cl) + s.cl.Alloc = s.mem +} + +func (s *FlightSqliteServerSuite) TearDownTest() { + s.Require().NoError(s.cl.Close()) + s.s.Shutdown() + s.srv = nil + s.mem.AssertSize(s.T(), 0) +} + +func (s *FlightSqliteServerSuite) fromJSON(dt arrow.DataType, json string) arrow.Array { + arr, _, _ := array.FromJSON(s.mem, dt, strings.NewReader(json)) + return arr +} + +func (s *FlightSqliteServerSuite) execCountQuery(query string) int64 { + info, err := s.cl.Execute(context.Background(), query) + s.NoError(err) + + rdr, err := s.cl.DoGet(context.Background(), info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + rec, err := rdr.Read() + s.NoError(err) + return rec.Column(0).(*array.Int64).Value(0) +} + +func (s *FlightSqliteServerSuite) TestCommandStatementQuery() { + ctx := context.Background() + info, err := s.cl.Execute(ctx, "SELECT * FROM intTable") + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + }, nil) + + s.Truef(expectedSchema.Equal(rec.Schema()), "expected: %s\ngot: %s", expectedSchema, rec.Schema()) + + idarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) + defer idarr.Release() + keyarr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) + defer keyarr.Release() + valarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) + defer valarr.Release() + foreignarr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) + defer foreignarr.Release() + + expectedRec := array.NewRecord(expectedSchema, []arrow.Array{idarr, keyarr, valarr, foreignarr}, 4) + defer expectedRec.Release() + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTables() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{}) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer catalogName.Release() + schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer schemaName.Release() + + tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) + defer tableName.Release() + + tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) + defer tableType.Release() + + expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableNameFilterPattern: proto.String("int%"), + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) + tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + expected := array.NewRecord(schema_ref.Tables, []arrow.Array{catalog, schema, table, tabletype}, 1) + defer func() { + catalog.Release() + schema.Release() + table.Release() + tabletype.Release() + expected.Release() + }() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + s.NoError(rdr.Err()) + + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithTableTypesFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableTypes: []string{"index"}, + }) + s.NoError(err) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(schema_ref.Tables.Equal(rdr.Schema()), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithExistingTableTypeFilter() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableTypes: []string{"table"}, + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalogName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer catalogName.Release() + schemaName := scalar.MakeArrayOfNull(arrow.BinaryTypes.String, 3, s.mem) + defer schemaName.Release() + + tableName := s.fromJSON(arrow.BinaryTypes.String, `["foreignTable", "intTable", "sqlite_sequence"]`) + defer tableName.Release() + + tableType := s.fromJSON(arrow.BinaryTypes.String, `["table", "table", "table"]`) + defer tableType.Release() + + expectedRec := array.NewRecord(schema_ref.Tables, []arrow.Array{catalogName, schemaName, tableName, tableType}, 3) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expectedRec, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTablesWithIncludedSchemas() { + ctx := context.Background() + info, err := s.cl.GetTables(ctx, &flightsql.GetTablesOpts{ + TableNameFilterPattern: proto.String("int%"), + IncludeSchema: true, + }) + s.NoError(err) + s.NotNil(info) + + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + catalog := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + schema := s.fromJSON(arrow.BinaryTypes.String, `[null]`) + table := s.fromJSON(arrow.BinaryTypes.String, `["intTable"]`) + tabletype := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + + dbTableName := "intTable" + + tableSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + {Name: "keyName", Type: arrow.BinaryTypes.String, + Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, dbTableName)}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, + Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, dbTableName)}, + }, nil) + schemaBuf := flight.SerializeSchema(tableSchema, s.mem) + binaryBldr := array.NewBinaryBuilder(s.mem, arrow.BinaryTypes.Binary) + binaryBldr.Append(schemaBuf) + schemaCol := binaryBldr.NewArray() + + expected := array.NewRecord(schema_ref.TablesWithIncludedSchema, []arrow.Array{catalog, schema, table, tabletype, schemaCol}, 1) + defer func() { + catalog.Release() + schema.Release() + table.Release() + tabletype.Release() + binaryBldr.Release() + schemaCol.Release() + expected.Release() + }() + + s.True(rdr.Next()) + rec := rdr.Record() + s.NotNil(rec) + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + s.NoError(rdr.Err()) + + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTypeInfo() { + ctx := context.Background() + info, err := s.cl.GetXdbcTypeInfo(ctx, nil) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := example.GetTypeInfoResult(s.mem) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTypeInfoFiltered() { + ctx := context.Background() + info, err := s.cl.GetXdbcTypeInfo(ctx, proto.Int32(-4)) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := example.GetFilteredTypeInfoResult(s.mem, -4) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetCatalogs() { + ctx := context.Background() + info, err := s.cl.GetCatalogs(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Schema().Equal(schema_ref.Catalogs), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetDbSchemas() { + ctx := context.Background() + info, err := s.cl.GetDBSchemas(ctx, &flightsql.GetDBSchemasOpts{}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Schema().Equal(schema_ref.DBSchemas), rdr.Schema().String()) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetTableTypes() { + ctx := context.Background() + info, err := s.cl.GetTableTypes(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + expected := s.fromJSON(arrow.BinaryTypes.String, `["table"]`) + defer expected.Release() + expectedRec := array.NewRecord(schema_ref.TableTypes, []arrow.Array{expected}, 1) + defer expectedRec.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expectedRec, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandStatementUpdate() { + ctx := context.Background() + result, err := s.cl.ExecuteUpdate(ctx, `INSERT INTO intTable (keyName, value) VALUES + ('KEYNAME1', 1001), ('KEYNAME2', 1002), ('KEYNAME3', 1003)`) + s.NoError(err) + s.EqualValues(3, result) + + result, err = s.cl.ExecuteUpdate(ctx, `UPDATE intTable SET keyName = 'KEYNAME1' + WHERE keyName = 'KEYNAME2' OR keyName = 'KEYNAME3'`) + s.NoError(err) + s.EqualValues(2, result) + + result, err = s.cl.ExecuteUpdate(ctx, `DELETE FROM intTable WHERE keyName = 'KEYNAME1'`) + s.NoError(err) + s.EqualValues(3, result) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQuery() { + ctx := context.Background() + prep, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable") + s.NoError(err) + defer prep.Close(ctx) + + info, err := prep.Execute(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) + + idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 2, 3, 4]`) + defer idArr.Release() + keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "zero", "negative one", null]`) + defer keyNameArr.Release() + valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 0, -1, null]`) + defer valueArr.Release() + foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1, 1, null]`) + defer foreignIdArr.Release() + + expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 4) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementQueryWithParams() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "SELECT * FROM intTable WHERE keyName LIKE ?") + s.NoError(err) + defer stmt.Close(ctx) + + typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[0]") + offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") + strArray := s.fromJSON(arrow.BinaryTypes.String, `["%one"]`) + bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") + bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[]") + dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") + paramArr, _ := array.NewDenseUnionFromArraysWithFields(typeIDs, + offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, + []string{"string", "bytes", "bigint", "double"}) + batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ + {Name: "parameter_1", Type: paramArr.DataType()}}, nil), + []arrow.Array{paramArr}, 1) + defer func() { + typeIDs.Release() + offsets.Release() + strArray.Release() + bytesArr.Release() + bigintArr.Release() + dblArr.Release() + paramArr.Release() + batch.Release() + }() + + stmt.SetParameters(batch) + info, err := stmt.Execute(ctx) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + + expectedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "keyName", Type: arrow.BinaryTypes.String, Metadata: s.getColMetadata(sqlite3.SQLITE_TEXT, ""), Nullable: true}, + {Name: "value", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}, + {Name: "foreignId", Type: arrow.PrimitiveTypes.Int64, Metadata: s.getColMetadata(sqlite3.SQLITE_INTEGER, ""), Nullable: true}}, nil) + + idArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 3]`) + defer idArr.Release() + keyNameArr := s.fromJSON(arrow.BinaryTypes.String, `["one", "negative one"]`) + defer keyNameArr.Release() + valueArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, -1]`) + defer valueArr.Release() + foreignIdArr := s.fromJSON(arrow.PrimitiveTypes.Int64, `[1, 1]`) + defer foreignIdArr.Release() + + expected := array.NewRecord(expectedSchema, []arrow.Array{idArr, keyNameArr, valueArr, foreignIdArr}, 2) + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdateWithParams() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', ?)") + s.NoError(err) + defer stmt.Close(ctx) + + typeIDs := s.fromJSON(arrow.PrimitiveTypes.Int8, "[2]") + offsets := s.fromJSON(arrow.PrimitiveTypes.Int32, "[0]") + strArray := s.fromJSON(arrow.BinaryTypes.String, "[]") + bytesArr := s.fromJSON(arrow.BinaryTypes.Binary, "[]") + bigintArr := s.fromJSON(arrow.PrimitiveTypes.Int64, "[999]") + dblArr := s.fromJSON(arrow.PrimitiveTypes.Float64, "[]") + paramArr, err := array.NewDenseUnionFromArraysWithFields(typeIDs, + offsets, []arrow.Array{strArray, bytesArr, bigintArr, dblArr}, + []string{"string", "bytes", "bigint", "double"}) + s.NoError(err) + batch := array.NewRecord(arrow.NewSchema([]arrow.Field{ + {Name: "parameter_1", Type: paramArr.DataType()}}, nil), + []arrow.Array{paramArr}, 1) + defer func() { + typeIDs.Release() + offsets.Release() + strArray.Release() + bytesArr.Release() + bigintArr.Release() + dblArr.Release() + paramArr.Release() + batch.Release() + }() + + stmt.SetParameters(batch) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + n, err := stmt.ExecuteUpdate(context.Background()) + s.NoError(err) + s.EqualValues(1, n) + s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + n, err = s.cl.ExecuteUpdate(context.Background(), "DELETE FROM intTable WHERE keyName = 'new_value'") + s.NoError(err) + s.EqualValues(1, n) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) +} + +func (s *FlightSqliteServerSuite) TestCommandPreparedStatementUpdate() { + ctx := context.Background() + stmt, err := s.cl.Prepare(ctx, s.mem, "INSERT INTO intTable (keyName, value) VALUES ('new_value', 999)") + s.NoError(err) + defer stmt.Close(ctx) + + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + result, err := stmt.ExecuteUpdate(ctx) + s.NoError(err) + s.EqualValues(1, result) + s.EqualValues(5, s.execCountQuery("SELECT COUNT(*) FROM intTable")) + result, err = s.cl.ExecuteUpdate(ctx, "DELETE FROM intTable WHERE keyName = 'new_value'") + s.NoError(err) + s.EqualValues(1, result) + s.EqualValues(4, s.execCountQuery("SELECT COUNT(*) FROM intTable")) +} + +func (s *FlightSqliteServerSuite) TestCommandGetPrimaryKeys() { + ctx := context.Background() + info, err := s.cl.GetPrimaryKeys(ctx, flightsql.TableRef{Table: "int%"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.PrimaryKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("intTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).(*array.Int32Builder).Append(1) + bldr.Field(5).AppendNull() + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetImportedKeys() { + ctx := context.Background() + info, err := s.cl.GetImportedKeys(ctx, flightsql.TableRef{Table: "intTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetExportedKeys() { + ctx := context.Background() + info, err := s.cl.GetExportedKeys(ctx, flightsql.TableRef{Table: "foreignTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func (s *FlightSqliteServerSuite) TestCommandGetCrossRef() { + ctx := context.Background() + info, err := s.cl.GetCrossReference(ctx, + flightsql.TableRef{Table: "foreignTable"}, + flightsql.TableRef{Table: "intTable"}) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + bldr := array.NewRecordBuilder(s.mem, schema_ref.ImportedKeys) + defer bldr.Release() + bldr.Field(0).AppendNull() + bldr.Field(1).AppendNull() + bldr.Field(2).(*array.StringBuilder).Append("foreignTable") + bldr.Field(3).(*array.StringBuilder).Append("id") + bldr.Field(4).AppendNull() + bldr.Field(5).AppendNull() + bldr.Field(6).(*array.StringBuilder).Append("intTable") + bldr.Field(7).(*array.StringBuilder).Append("foreignId") + bldr.Field(8).(*array.Int32Builder).Append(0) + bldr.Field(9).AppendNull() + bldr.Field(10).AppendNull() + bldr.Field(11).(*array.Uint8Builder).Append(3) + bldr.Field(12).(*array.Uint8Builder).Append(3) + expected := bldr.NewRecord() + defer expected.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + s.Truef(array.RecordEqual(expected, rec), "expected: %s\ngot: %s", expected, rec) + s.False(rdr.Next()) +} + +func validateSqlInfo(t *testing.T, expected interface{}, sc scalar.Scalar) bool { + switch ex := expected.(type) { + case string: + return assert.Equal(t, ex, sc.String()) + case bool: + return assert.Equal(t, ex, sc.(*scalar.Boolean).Value) + case int64: + return assert.Equal(t, ex, sc.(*scalar.Int64).Value) + case int32: + return assert.Equal(t, ex, sc.(*scalar.Int32).Value) + case []string: + arr := sc.(*scalar.List).Value.(*array.String) + assert.EqualValues(t, len(ex), arr.Len()) + for i, v := range ex { + assert.Equal(t, v, arr.Value(i)) + } + case map[int32][]int32: + // map is a list of structs with key and values + structArr := sc.(*scalar.Map).Value.(*array.Struct) + keys := structArr.Field(0).(*array.Int32) + values := structArr.Field(1).(*array.List) + // assert that the map has the right size + assert.EqualValues(t, len(ex), keys.Len()) + + // for each element, match the argument + for i := 0; i < keys.Len(); i++ { + keyScalar, _ := scalar.GetScalar(keys, i) + infoID := keyScalar.(*scalar.Int32).Value + + // assert the key exists + list, ok := ex[infoID] + assert.True(t, ok) + + // assert the int32list is the right size + start, end := values.ValueOffsets(i) + assert.EqualValues(t, len(list), end-start) + + // for each element make sure it matches + for j, v := range list { + listItem, err := scalar.GetScalar(values.ListValues(), int(start)+j) + assert.NoError(t, err) + assert.Equal(t, v, listItem.(*scalar.Int32).Value) + } + } + } + return true +} + +func (s *FlightSqliteServerSuite) TestCommandGetSqlInfo() { + expectedResults := example.SqlInfoResultMap() + infoIDs := make([]flightsql.SqlInfo, 0, len(expectedResults)) + for k := range expectedResults { + infoIDs = append(infoIDs, flightsql.SqlInfo(k)) + } + + ctx := context.Background() + info, err := s.cl.GetSqlInfo(ctx, infoIDs) + s.NoError(err) + rdr, err := s.cl.DoGet(ctx, info.Endpoint[0].Ticket) + s.NoError(err) + defer rdr.Release() + + s.True(rdr.Next()) + rec := rdr.Record() + rec.Retain() + defer rec.Release() + s.False(rdr.Next()) + + s.EqualValues(2, rec.NumCols()) + s.EqualValues(len(expectedResults), rec.NumRows()) + + colName := rec.Column(0).(*array.Uint32) + colValue := rec.Column(1) + for i := 0; i < int(rec.NumRows()); i++ { + expected := expectedResults[colName.Value(i)] + sc, err := scalar.GetScalar(colValue, i) + s.NoError(err) + + s.True(validateSqlInfo(s.T(), expected, sc.(*scalar.DenseUnion).ChildValue())) + + sc.(*scalar.DenseUnion).Release() + } +} + +func TestSqliteServer(t *testing.T) { + suite.Run(t, new(FlightSqliteServerSuite)) +} From 68df566697ad6b082c3471c992642758bd5ba7e6 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 11:45:38 -0400 Subject: [PATCH 17/23] updates from feedback --- dev/release/post-11-bump-versions-test.rb | 5 +++++ dev/release/utils-prepare.sh | 3 +++ go/arrow/array/list_test.go | 2 ++ go/arrow/doc.go | 2 ++ go/arrow/flight/flightsql/client.go | 4 ---- go/arrow/flight/flightsql/example/sqlite_info.go | 7 +++++-- go/arrow/flight/flightsql/schema_ref/reference_schemas.go | 2 +- 7 files changed, 18 insertions(+), 7 deletions(-) diff --git a/dev/release/post-11-bump-versions-test.rb b/dev/release/post-11-bump-versions-test.rb index 6770ca1c22edd..2486bf4aacb4f 100644 --- a/dev/release/post-11-bump-versions-test.rb +++ b/dev/release/post-11-bump-versions-test.rb @@ -206,6 +206,11 @@ def test_version_post_tag "-\tDefaultCreatedBy = \"parquet-go version #{@snapshot_version}\"", "+\tDefaultCreatedBy = \"parquet-go version #{@next_snapshot_version}\"", ] + else if path == "go/arrow/doc.go" + hunks << [ + "-const PkgVersion = \"#{@snapshot_version}\"", + "+const PkgVersion = \"#{@next_snapshot_version}\"", + ] end expected_changes << {hunks: hunks, path: path} end diff --git a/dev/release/utils-prepare.sh b/dev/release/utils-prepare.sh index c68632b9d0487..49af5c608aa1a 100644 --- a/dev/release/utils-prepare.sh +++ b/dev/release/utils-prepare.sh @@ -156,6 +156,9 @@ update_versions() { sed -i.bak -E -e \ "s/\"parquet-go version .+\"/\"parquet-go version ${version}\"/" \ parquet/writer_properties.go + sed -i.bak -E -e \ + "s/const PkgVersion = \".*/const PkgVersion = \"${version}\"/" \ + arrow/doc.go find . -name "*.bak" -exec rm {} \; git add . popd diff --git a/go/arrow/array/list_test.go b/go/arrow/array/list_test.go index f493167f76a1c..eb09f655d7e52 100644 --- a/go/arrow/array/list_test.go +++ b/go/arrow/array/list_test.go @@ -33,6 +33,8 @@ func TestListArray(t *testing.T) { }{ {arrow.LIST, []int32{0, 3, 3, 3, 7}, arrow.ListOf(arrow.PrimitiveTypes.Int32)}, {arrow.LARGE_LIST, []int64{0, 3, 3, 3, 7}, arrow.LargeListOf(arrow.PrimitiveTypes.Int32)}, + {arrow.LIST, []int32{0, 3, 3, 3, 7}, arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.PrimitiveTypes.Int32, Nullable: true})}, + {arrow.LARGE_LIST, []int64{0, 3, 3, 3, 7}, arrow.LargeListOfField(arrow.Field{Name: "item", Type: arrow.PrimitiveTypes.Int32, Nullable: true})}, } for _, tt := range tests { diff --git a/go/arrow/doc.go b/go/arrow/doc.go index 0af5cd163ab1b..cf73f1a00b3f7 100644 --- a/go/arrow/doc.go +++ b/go/arrow/doc.go @@ -31,6 +31,8 @@ array is valid (not null). If the array has no null entries, it is possible to o */ package arrow +const PkgVersion = "10.0.0-SNAPSHOT" + //go:generate go run _tools/tmpl/main.go -i -data=numeric.tmpldata type_traits_numeric.gen.go.tmpl type_traits_numeric.gen_test.go.tmpl array/numeric.gen.go.tmpl array/numericbuilder.gen.go.tmpl array/bufferbuilder_numeric.gen.go.tmpl //go:generate go run _tools/tmpl/main.go -i -data=datatype_numeric.gen.go.tmpldata datatype_numeric.gen.go.tmpl tensor/numeric.gen.go.tmpl tensor/numeric.gen_test.go.tmpl //go:generate go run _tools/tmpl/main.go -i -data=scalar/numeric.gen.go.tmpldata scalar/numeric.gen.go.tmpl scalar/numeric.gen_test.go.tmpl diff --git a/go/arrow/flight/flightsql/client.go b/go/arrow/flight/flightsql/client.go index 077f5295c9f4c..5f7f693d2b2f7 100644 --- a/go/arrow/flight/flightsql/client.go +++ b/go/arrow/flight/flightsql/client.go @@ -232,10 +232,6 @@ func (c *Client) GetSqlInfo(ctx context.Context, info []SqlInfo, opts ...grpc.Ca return flightInfoForCommand(ctx, c, cmd, opts...) } -// Prepare creates a PreparedStatement object for the specified query. -// The resulting PreparedStatement object should be Closed when no longer -// needed. It will maintain a reference to this Client for use to execute -// and use the specified allocator for any allocations it needs to perform. // Prepare creates a PreparedStatement object for the specified query. // The resulting PreparedStatement object should be Closed when no longer // needed. It will maintain a reference to this Client for use to execute diff --git a/go/arrow/flight/flightsql/example/sqlite_info.go b/go/arrow/flight/flightsql/example/sqlite_info.go index 5a85b21412866..e4dcd160b0aca 100644 --- a/go/arrow/flight/flightsql/example/sqlite_info.go +++ b/go/arrow/flight/flightsql/example/sqlite_info.go @@ -19,13 +19,16 @@ package example -import "github.com/apache/arrow/go/v10/arrow/flight/flightsql" +import ( + "github.com/apache/arrow/go/v10/arrow" + "github.com/apache/arrow/go/v10/arrow/flight/flightsql" +) func SqlInfoResultMap() flightsql.SqlInfoResultMap { return flightsql.SqlInfoResultMap{ uint32(flightsql.SqlInfoFlightSqlServerName): "db_name", uint32(flightsql.SqlInfoFlightSqlServerVersion): "sqlite 3", - uint32(flightsql.SqlInfoFlightSqlServerArrowVersion): "10.0.0-SNAPSHOT", + uint32(flightsql.SqlInfoFlightSqlServerArrowVersion): arrow.PkgVersion, uint32(flightsql.SqlInfoFlightSqlServerReadOnly): false, uint32(flightsql.SqlInfoDDLCatalog): false, uint32(flightsql.SqlInfoDDLSchema): false, diff --git a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go index 728280b5dee01..9f9914d9f8a41 100644 --- a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go +++ b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go @@ -64,7 +64,7 @@ var ( {Name: "fk_key_name", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "pk_key_name", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "update_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, - {Name: "delete_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: true}, + {Name: "delete_rule", Type: arrow.PrimitiveTypes.Uint8, Nullable: false}, }, nil) ImportedKeys = ImportedExportedKeysAndCrossReference ExportedKeys = ImportedExportedKeysAndCrossReference From c2a9a6a5c295301e888a59c0ec0079f053eea819 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 11:55:34 -0400 Subject: [PATCH 18/23] fix last nullability for c++ comparison --- go/arrow/flight/flightsql/schema_ref/reference_schemas.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go index 9f9914d9f8a41..7a4a14064d540 100644 --- a/go/arrow/flight/flightsql/schema_ref/reference_schemas.go +++ b/go/arrow/flight/flightsql/schema_ref/reference_schemas.go @@ -88,7 +88,7 @@ var ( {Name: "column_size", Type: arrow.PrimitiveTypes.Int32, Nullable: true}, {Name: "literal_prefix", Type: arrow.BinaryTypes.String, Nullable: true}, {Name: "literal_suffix", Type: arrow.BinaryTypes.String, Nullable: true}, - {Name: "create_params", Type: arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.BinaryTypes.String, Nullable: false})}, + {Name: "create_params", Type: arrow.ListOfField(arrow.Field{Name: "item", Type: arrow.BinaryTypes.String, Nullable: false}), Nullable: true}, {Name: "nullable", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, {Name: "case_sensitive", Type: arrow.FixedWidthTypes.Boolean, Nullable: false}, {Name: "searchable", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, From f32c2f131218c16d03f14693b59d758bfe7ed23a Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 13:01:50 -0400 Subject: [PATCH 19/23] fix ruby version test --- dev/release/post-11-bump-versions-test.rb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/release/post-11-bump-versions-test.rb b/dev/release/post-11-bump-versions-test.rb index 2486bf4aacb4f..4aafe4cf82259 100644 --- a/dev/release/post-11-bump-versions-test.rb +++ b/dev/release/post-11-bump-versions-test.rb @@ -206,7 +206,7 @@ def test_version_post_tag "-\tDefaultCreatedBy = \"parquet-go version #{@snapshot_version}\"", "+\tDefaultCreatedBy = \"parquet-go version #{@next_snapshot_version}\"", ] - else if path == "go/arrow/doc.go" + elsif path == "go/arrow/doc.go" hunks << [ "-const PkgVersion = \"#{@snapshot_version}\"", "+const PkgVersion = \"#{@next_snapshot_version}\"", From ef422925b20932b94c8c3c1c98a278c732e49cc4 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 13:42:19 -0400 Subject: [PATCH 20/23] explicitly downgrade zlib to fix js integration? --- ci/docker/conda-integration.dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ci/docker/conda-integration.dockerfile b/ci/docker/conda-integration.dockerfile index 8bcf5954d1db7..a512ae7d6e5eb 100644 --- a/ci/docker/conda-integration.dockerfile +++ b/ci/docker/conda-integration.dockerfile @@ -35,7 +35,8 @@ RUN mamba install -q -y \ maven=${maven} \ nodejs=${node} \ yarn \ - openjdk=${jdk} && \ + openjdk=${jdk} \ + zlib=1.2.11 && \ mamba clean --all --force-pkgs-dirs # Install Rust with only the needed components From db7a6a6955f5b09eea5eabe72c5fc96143db618f Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 16:00:58 -0400 Subject: [PATCH 21/23] dev version bump tests --- dev/release/01-prepare-test.rb | 7 +++++++ dev/release/post-11-bump-versions-test.rb | 20 +++++++++++++------- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/dev/release/01-prepare-test.rb b/dev/release/01-prepare-test.rb index b498a29763241..d0b47ec430e18 100644 --- a/dev/release/01-prepare-test.rb +++ b/dev/release/01-prepare-test.rb @@ -169,6 +169,13 @@ def test_version_pre_tag ], ], }, + { + path: "go/arrow/doc.go", + hunks: [ + ["-const PkgVersion = \"#{@snapshot_version}\"", + "+const PkgVersion = \"#{@release_version}\""], + ], + }. { path: "go/parquet/writer_properties.go", hunks: [ diff --git a/dev/release/post-11-bump-versions-test.rb b/dev/release/post-11-bump-versions-test.rb index 4aafe4cf82259..01e063a99bfe8 100644 --- a/dev/release/post-11-bump-versions-test.rb +++ b/dev/release/post-11-bump-versions-test.rb @@ -185,6 +185,17 @@ def test_version_post_tag ] Dir.glob("go/**/{go.mod,*.go,*.go.*}") do |path| + if path == "go/arrow/doc.go" + expected_changes << { + path: path, + hunks: [ + [ + "-const PkgVersion = \"#{@snapshot_version}\"", + "+const PkgVersion = \"#{@next_snapshot_version}\"", + ] + ]} + next + end import_path = "github.com/apache/arrow/go/v#{@snapshot_major_version}" lines = File.readlines(path, chomp: true) target_lines = lines.grep(/#{Regexp.escape(import_path)}/) @@ -205,15 +216,10 @@ def test_version_post_tag hunks << [ "-\tDefaultCreatedBy = \"parquet-go version #{@snapshot_version}\"", "+\tDefaultCreatedBy = \"parquet-go version #{@next_snapshot_version}\"", - ] - elsif path == "go/arrow/doc.go" - hunks << [ - "-const PkgVersion = \"#{@snapshot_version}\"", - "+const PkgVersion = \"#{@next_snapshot_version}\"", - ] + ] end expected_changes << {hunks: hunks, path: path} - end + end Dir.glob("java/**/pom.xml") do |path| version = "#{@snapshot_version}" From 29966333c1c1bcb2f50e45053176344ab057e8b3 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 16:03:57 -0400 Subject: [PATCH 22/23] gah put a silly period instead of a comma --- dev/release/01-prepare-test.rb | 2 +- dev/release/post-11-bump-versions-test.rb | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/release/01-prepare-test.rb b/dev/release/01-prepare-test.rb index d0b47ec430e18..fa00d6290c1f0 100644 --- a/dev/release/01-prepare-test.rb +++ b/dev/release/01-prepare-test.rb @@ -175,7 +175,7 @@ def test_version_pre_tag ["-const PkgVersion = \"#{@snapshot_version}\"", "+const PkgVersion = \"#{@release_version}\""], ], - }. + }, { path: "go/parquet/writer_properties.go", hunks: [ diff --git a/dev/release/post-11-bump-versions-test.rb b/dev/release/post-11-bump-versions-test.rb index 01e063a99bfe8..51eb4bb3077aa 100644 --- a/dev/release/post-11-bump-versions-test.rb +++ b/dev/release/post-11-bump-versions-test.rb @@ -192,7 +192,7 @@ def test_version_post_tag [ "-const PkgVersion = \"#{@snapshot_version}\"", "+const PkgVersion = \"#{@next_snapshot_version}\"", - ] + ], ]} next end From 72b75b141c22ca1aebec37346cdb95ec4e81cc0b Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 15 Aug 2022 17:09:55 -0400 Subject: [PATCH 23/23] remove trailing whitespace --- dev/release/post-11-bump-versions-test.rb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/release/post-11-bump-versions-test.rb b/dev/release/post-11-bump-versions-test.rb index 51eb4bb3077aa..1c87a9ea45da2 100644 --- a/dev/release/post-11-bump-versions-test.rb +++ b/dev/release/post-11-bump-versions-test.rb @@ -216,10 +216,10 @@ def test_version_post_tag hunks << [ "-\tDefaultCreatedBy = \"parquet-go version #{@snapshot_version}\"", "+\tDefaultCreatedBy = \"parquet-go version #{@next_snapshot_version}\"", - ] + ] end expected_changes << {hunks: hunks, path: path} - end + end Dir.glob("java/**/pom.xml") do |path| version = "#{@snapshot_version}"