joellubi commented on code in PR #38385:
URL: https://github.com/apache/arrow/pull/38385#discussion_r1562952457
##########
go/arrow/internal/flight_integration/scenario.go:
##########
@@ -2948,3 +2951,136 @@ func (tester *sessionOptionsScenarioTester)
ValidateSeventhGetSessionOptions(ctx
return nil
}
+
+type flightSqlIngestionScenarioTester struct {
+ flightsql.BaseServer
+}
+
+func (m *flightSqlIngestionScenarioTester) MakeServer(port int) flight.Server {
+ srv := flight.NewServerWithMiddleware(nil)
+ m.RegisterSqlInfo(flightsql.SqlInfoFlightSqlServerBulkIngestion, true)
+
m.RegisterSqlInfo(flightsql.SqlInfoFlightSqlServerIngestTransactionsSupported,
true)
+
+ srv.RegisterFlightService(flightsql.NewFlightServer(m))
+ initServer(port, srv)
+ return srv
+}
+
+func (m *flightSqlIngestionScenarioTester) RunClient(addr string, opts
...grpc.DialOption) error {
+ client, err := flightsql.NewClient(addr, nil, nil, opts...)
+ if err != nil {
+ return err
+ }
+ defer client.Close()
+
+ return m.ValidateIngestion(client)
+}
+
+func (m *flightSqlIngestionScenarioTester) ValidateIngestion(client
*flightsql.Client) error {
+ ctx := context.Background()
+ opts := getIngestOptions()
+ ingestResult, err := client.ExecuteIngest(ctx, getIngestRecords(),
&opts)
+ if err != nil {
+ return err
+ }
+ if ingestResult != ingestStatementExpectedRows {
+ return fmt.Errorf("expected ingest return %d got %d",
ingestStatementExpectedRows, ingestResult)
+ }
+ return nil
+}
+
+func (m *flightSqlIngestionScenarioTester) DoPutCommandStatementIngest(ctx
context.Context, cmd flightsql.StatementIngest, rdr flight.MessageReader)
(int64, error) {
+ expectedSchema := getIngestSchema()
+ expectedOpts := getIngestOptions()
+
+ if err := assertEq(expectedOpts.TableDefinitionOptions.IfExists,
cmd.GetTableDefinitionOptions().IfExists); err != nil {
+ return 0, err
+ }
+
+ if err := assertEq(expectedOpts.TableDefinitionOptions.IfNotExist,
cmd.GetTableDefinitionOptions().IfNotExist); err != nil {
+ return 0, err
+ }
+
+ if err := assertEq(expectedOpts.Table, cmd.GetTable()); err != nil {
+ return 0, err
+ }
+
+ if err := assertEq(*expectedOpts.Schema, cmd.GetSchema()); err != nil {
+ return 0, err
+ }
+
+ if err := assertEq(*expectedOpts.Catalog, cmd.GetCatalog()); err != nil
{
+ return 0, err
+ }
+
+ if err := assertEq(expectedOpts.Temporary, cmd.GetTemporary()); err !=
nil {
+ return 0, err
+ }
+
+ if err := assertEq(expectedOpts.TransactionId, cmd.GetTransactionId());
err != nil {
+ return 0, err
+ }
+
+ if err := assertEq(expectedOpts.Options, cmd.GetOptions()); err != nil {
+ return 0, err
+ }
+
+ var nRecords int64
+ for rdr.Next() {
+ rec := rdr.Record()
+ nRecords += rec.NumRows()
+
+ if err := assertEq(true, expectedSchema.Equal(rec.Schema()));
err != nil {
+ return 0, err
+ }
+ }
+
+ return nRecords, nil
+}
+
+// Options to assert before/after mocked ingest call
+func getIngestOptions() flightsql.ExecuteIngestOpts {
Review Comment:
Updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
