| 636 | } |
| 637 | |
| 638 | Status FlightSqlServerBase::GetSchema(const ServerCallContext& context, |
| 639 | const FlightDescriptor& request, |
| 640 | std::unique_ptr<SchemaResult>* schema) { |
| 641 | google::protobuf::Any any; |
| 642 | if (!any.ParseFromArray(request.cmd.data(), static_cast<int>(request.cmd.size()))) { |
| 643 | return Status::Invalid("Unable to parse command"); |
| 644 | } |
| 645 | |
| 646 | if (any.Is<pb::sql::CommandStatementQuery>()) { |
| 647 | ARROW_ASSIGN_OR_RAISE(StatementQuery internal_command, |
| 648 | ParseCommandStatementQuery(any)); |
| 649 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 650 | GetSchemaStatement(context, internal_command, request)); |
| 651 | return Status::OK(); |
| 652 | } else if (any.Is<pb::sql::CommandStatementSubstraitPlan>()) { |
| 653 | ARROW_ASSIGN_OR_RAISE(StatementSubstraitPlan internal_command, |
| 654 | ParseCommandStatementSubstraitPlan(any)); |
| 655 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 656 | GetSchemaSubstraitPlan(context, internal_command, request)); |
| 657 | return Status::OK(); |
| 658 | } else if (any.Is<pb::sql::CommandPreparedStatementQuery>()) { |
| 659 | ARROW_ASSIGN_OR_RAISE(PreparedStatementQuery internal_command, |
| 660 | ParseCommandPreparedStatementQuery(any)); |
| 661 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 662 | GetSchemaPreparedStatement(context, internal_command, request)); |
| 663 | return Status::OK(); |
| 664 | } else if (any.Is<pb::sql::CommandGetCatalogs>()) { |
| 665 | ARROW_ASSIGN_OR_RAISE(*schema, SchemaResult::Make(*SqlSchema::GetCatalogsSchema())); |
| 666 | return Status::OK(); |
| 667 | } else if (any.Is<pb::sql::CommandGetCrossReference>()) { |
| 668 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 669 | SchemaResult::Make(*SqlSchema::GetCrossReferenceSchema())); |
| 670 | return Status::OK(); |
| 671 | } else if (any.Is<pb::sql::CommandGetDbSchemas>()) { |
| 672 | ARROW_ASSIGN_OR_RAISE(*schema, SchemaResult::Make(*SqlSchema::GetDbSchemasSchema())); |
| 673 | return Status::OK(); |
| 674 | } else if (any.Is<pb::sql::CommandGetExportedKeys>()) { |
| 675 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 676 | SchemaResult::Make(*SqlSchema::GetExportedKeysSchema())); |
| 677 | return Status::OK(); |
| 678 | } else if (any.Is<pb::sql::CommandGetImportedKeys>()) { |
| 679 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 680 | SchemaResult::Make(*SqlSchema::GetImportedKeysSchema())); |
| 681 | return Status::OK(); |
| 682 | } else if (any.Is<pb::sql::CommandGetPrimaryKeys>()) { |
| 683 | ARROW_ASSIGN_OR_RAISE(*schema, |
| 684 | SchemaResult::Make(*SqlSchema::GetPrimaryKeysSchema())); |
| 685 | return Status::OK(); |
| 686 | } else if (any.Is<pb::sql::CommandGetSqlInfo>()) { |
| 687 | ARROW_ASSIGN_OR_RAISE(*schema, SchemaResult::Make(*SqlSchema::GetSqlInfoSchema())); |
| 688 | return Status::OK(); |
| 689 | } else if (any.Is<pb::sql::CommandGetTables>()) { |
| 690 | ARROW_ASSIGN_OR_RAISE(GetTables command, ParseCommandGetTables(any)); |
| 691 | if (command.include_schema) { |
| 692 | ARROW_ASSIGN_OR_RAISE( |
| 693 | *schema, SchemaResult::Make(*SqlSchema::GetTablesSchemaWithIncludedSchema())); |
| 694 | } else { |
| 695 | ARROW_ASSIGN_OR_RAISE(*schema, SchemaResult::Make(*SqlSchema::GetTablesSchema())); |