diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_foreign_keys.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_foreign_keys.cc index c557c9716f..6257a2d4fc 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_foreign_keys.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_foreign_keys.cc @@ -13,11 +13,17 @@ // limitations under the License. #include "google/cloud/odbc/bq_driver/internal/odbc_sql_foreign_keys.h" +#include "google/cloud/odbc/bq_driver/internal/odbc_sql_columns.h" #include "google/cloud/odbc/bq_driver/internal/trace_utils.h" +#include "absl/strings/match.h" +#include +#include #include namespace google::cloud::odbc_bq_driver_internal { +using ::google::cloud::bigquery_v2_minimal_internal::ColumnReference; +using ::google::cloud::bigquery_v2_minimal_internal::ForeignKey; using ::google::cloud::odbc_internal::SQLStates; using ::google::cloud::odbc_internal::StatusRecord; using ::google::cloud::odbc_internal::StatusRecordOr; @@ -80,6 +86,120 @@ std::string const kBasicForeignKeysQueryPrefix = std::string const kBasicForeignKeysQuerySuffix = "ORDER BY pk_table, pk_column_ordinal_position, pk_column"; +// BigQuery INFORMATION_SCHEMA reports constraint names as +// ".", where the primary key constraint name is +// always "pk$" (BigQuery primary keys cannot be named) and an unnamed foreign +// key constraint is named "fk$" by order of declaration, which is the order +// tables.get lists them in. These suffixes let the tables.get based path below +// produce the same constraint names as the INFORMATION_SCHEMA query path. +std::string const kPrimaryKeyNameSuffix = ".pk$"; +std::string const kForeignKeyNameSuffix = ".fk$"; + +// Fetches the foreign keys of a single, exactly named table by reading the +// table's constraints via the tables.get REST API, the same way +// FetchPKResultSetFromTableMetaData does for SQLPrimaryKeys. This avoids the +// several seconds of fixed latency an INFORMATION_SCHEMA query job carries +// (table names are ordinary identifiers, not search patterns, per the ODBC +// spec). When pk_table_name is non-empty, only foreign keys referencing that +// table are returned. The rows are returned as a ready ResultSet inside +// DSResults, which ProcessQueryResults passes through unchanged. +StatusRecordOr FetchForeignKeysFromTableMetadata( + StatementHandle& stmt_handle, std::string const& catalog_name, + std::string const& schema_name, std::string const& pk_table_name, + std::string const& fk_table_name) { + ConnectionHandle& conn_handle = *(stmt_handle.GetConnectionHandle()); + auto bq_table_status = + FetchBQTableData(conn_handle, catalog_name, schema_name, fk_table_name); + ResultSet result_set; + result_set.row_schema.resize(kForeignKeysMap.size()); + for (auto const& [_, schema] : kForeignKeysMap) { + result_set.row_schema[schema.col_index] = schema; + } + if (!bq_table_status) { + auto const& status = bq_table_status.GetStatusRecord(); + // An unknown table produces an empty result set, matching the + // INFORMATION_SCHEMA query path. + if (status.native_error_code == 404) { + LOG(INFO) << "FetchForeignKeysFromTableMetadata:: Table not found: '" + << catalog_name << "." << schema_name << "." << fk_table_name + << "'"; + DSResults ds_results; + ds_results.data_source_results = std::move(result_set); + return ds_results; + } + LOG(ERROR) << "FetchForeignKeysFromTableMetadata::FetchBQTableData:: " + << status.message; + stmt_handle.GetDiagnostics().AddStatusRecord(status); + return status; + } + // Ordinal of the foreign key within the table's constraints, which is what + // BigQuery numbers the unnamed ones by. Counted over all of them, not just + // the ones kept by the pk_table_name filter below. + int fk_ordinal = 0; + for (ForeignKey const& foreign_key : + bq_table_status->table_constraints.foreign_keys) { + ++fk_ordinal; + if (!pk_table_name.empty() && + foreign_key.referenced_table.table_id != pk_table_name) { + continue; + } + // INFORMATION_SCHEMA reports constraint names prefixed with the table id, + // e.g. "my_table.fk$1" / "my_table.my_named_fk"; reproduce that. The name + // of an explicitly named constraint is not available here: tables.get + // reports it in "name", but google-cloud-cpp parses it from "keyName", so + // ForeignKey::key_name always arrives empty and such a constraint is + // reported as "fk$" instead of its declared name. key_name is still + // preferred when present, so this corrects itself if the dependency does. + DSValue fk_name_value; + if (foreign_key.key_name.empty()) { + StringToDSValue( + fk_table_name + kForeignKeyNameSuffix + std::to_string(fk_ordinal), + fk_name_value); + } else { + StringToDSValue(fk_table_name + "." + foreign_key.key_name, + fk_name_value); + } + DSValue pk_name_value; + StringToDSValue( + foreign_key.referenced_table.table_id + kPrimaryKeyNameSuffix, + pk_name_value); + SQLBIGINT key_seq = 0; + for (ColumnReference const& column_ref : foreign_key.column_references) { + ++key_seq; + DSRow row(kForeignKeysMap.size()); + StringToDSValue(foreign_key.referenced_table.project_id, + row[0]); // PKTABLE_CAT + StringToDSValue(foreign_key.referenced_table.dataset_id, + row[1]); // PKTABLE_SCHEM + StringToDSValue(foreign_key.referenced_table.table_id, + row[2]); // PKTABLE_NAME + StringToDSValue(column_ref.referenced_column, row[3]); // PKCOLUMN_NAME + StringToDSValue(catalog_name, row[4]); // FKTABLE_CAT + StringToDSValue(schema_name, row[5]); // FKTABLE_SCHEM + StringToDSValue(fk_table_name, row[6]); // FKTABLE_NAME + StringToDSValue(column_ref.referencing_column, row[7]); // FKCOLUMN_NAME + ArithmeticToDSValue(key_seq, row[8]); // KEY_SEQ + row[9] = kNullValue; // UPDATE_RULE + row[10] = kNullValue; // DELETE_RULE + row[11] = fk_name_value; // FK_NAME + row[12] = pk_name_value; // PK_NAME + ArithmeticToDSValue(2, row[13]); // DEFERRABILITY + result_set.rows.push_back(std::move(row)); + } + } + // Match the query path ordering: PKTABLE_NAME, then key sequence. + std::stable_sort(result_set.rows.begin(), result_set.rows.end(), + [](DSRow const& a, DSRow const& b) { + if (a[2] != b[2]) return a[2] < b[2]; + DSValue seq_a = a[8]; + DSValue seq_b = b[8]; + return DSValueToInt(seq_a) < DSValueToInt(seq_b); + }); + DSResults ds_results; + ds_results.data_source_results = std::move(result_set); + return ds_results; +} + } // namespace odbc_internal::StatusRecordOr FetchForeignKeysFromDataSource( @@ -156,6 +276,17 @@ odbc_internal::StatusRecordOr FetchForeignKeysFromDataSource( stmt_handle.GetDiagnostics().AddStatusRecord(status_record); return status_record; } + // Fast path: the foreign key table is specified with an exact (non-pattern) + // name — read that table's constraints directly via the tables.get REST API + // instead of running an INFORMATION_SCHEMA query job. This covers both the + // FK-only and the PK+FK cases; the PK-only case cannot use it because the + // referencing tables are not known upfront. Names containing '%' fall back + // to the query below, which preserves the historical LIKE matching. + if (!fk_table_name.empty() && !absl::StrContains(fk_table_name, '%') && + !absl::StrContains(pk_table_name, '%')) { + return FetchForeignKeysFromTableMetadata( + stmt_handle, catalog_name, schema_name, pk_table_name, fk_table_name); + } // Construct named query for foreign keys. std::string foreign_keys_query(kBasicForeignKeysQueryPrefix); foreign_keys_query