Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ struct Oauth {
std::string kms_key_name;
std::string psc;
GCD gcd;
std::string impersonated_email;
};

// Returns true if all required BYOID properties are set.
Expand Down
5 changes: 5 additions & 0 deletions google/cloud/odbc/bq_client_interface/odbc_bq_client.cc
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,11 @@ StatusRecordOr<std::shared_ptr<ODBCBQClient>> ODBCBQClient::CreateBQClient(
return credentials.GetStatusRecord();
}

if (!oauth.impersonated_email.empty()) {
credentials = google::cloud::MakeImpersonateServiceAccountCredentials(
credentials.GetValue(), oauth.impersonated_email, options);
}

options.set<google::cloud::UnifiedCredentialsOption>(*credentials);

if (oauth.gcd.enable_gcd && oauth.gcd.universe_domain != "googleapis.com") {
Expand Down
2 changes: 2 additions & 0 deletions google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,8 @@ void ConnectionHandle::SetUp(Section& dsn_section,
attribute_str_values_.insert({SQL_ATTR_CURRENT_CATALOG, dsn_.catalog});
}

dsn_.impersonated_email = dsn_section["SERVICEACCOUNTIMPERSONATIONEMAIL"];

// Populate HTAPI related configurations
std::string use_default_large_results_dataset =
dsn_section["USEDEFAULTLARGERESULTSDATASET"];
Expand Down
1 change: 1 addition & 0 deletions google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ struct Dsn {
std::string psc;
bool enable_gcd;
std::string universe_domain;
std::string impersonated_email;
};

class EnvironmentHandle;
Expand Down
1 change: 1 addition & 0 deletions google/cloud/odbc/bq_driver/odbc_connection.cc
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ Authentication CreateAuth(Dsn const& dsn) {
auth.oauth.psc = dsn.psc;
auth.oauth.gcd.enable_gcd = dsn.enable_gcd;
auth.oauth.gcd.universe_domain = dsn.universe_domain;
auth.oauth.impersonated_email = dsn.impersonated_email;
return auth;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4365,6 +4365,121 @@ TEST(SQLMoreResults, ProcedureWithDescriptorAndQueryParams) {
table.Drop(conn);
EXPECT_EQ(Disconnect(conn), SQL_SUCCESS);
}
TEST(StatementTest, VerifyServiceAccountImpersonationEmail) {
struct TestCase {
std::string impersonation_email;
bool expect_success;
};

std::vector<TestCase> const test_cases = {
{
"bq-devtools-simba-drivers-test@bigquery-devtools-drivers."
"iam.gserviceaccount.com",
true,
},
{
"kirl-test@bigquery-devtools-drivers.iam.gserviceaccount.com",
false,
},
};

for (auto const& test : test_cases) {
SCOPED_TRACE(test.impersonation_email);

auto conn = std::make_shared<ODBCHandles>();

std::string conn_str =
kDefaultConnectionString +
";ServiceAccountImpersonationEmail=" + test.impersonation_email;

EXPECT_EQ(Connect(conn_str, conn), SQL_SUCCESS);

SQLRETURN status =
SQLExecDirect(conn->hstmt, (SQLCHAR*)"SELECT SESSION_USER()", SQL_NTS);

if (test.expect_success) {
EXPECT_EQ(status, SQL_SUCCESS);

status = SQLFetch(conn->hstmt);
EXPECT_EQ(status, SQL_SUCCESS);

SQLWCHAR buf[256] = {};
SQLLEN indicator = 0;
std::wstring result;

do {
memset(buf, 0, sizeof(buf));

status = SQLGetData(conn->hstmt, 1, SQL_C_WCHAR, buf, sizeof(buf),
&indicator);

CheckError(status, "SQLGetData", conn);

if (indicator == SQL_NULL_DATA) {
break;
}

for (size_t i = 0; i < sizeof(buf) / sizeof(SQLWCHAR) && buf[i] != 0;
++i) {
result.push_back(buf[i]);
}

} while (status == SQL_SUCCESS_WITH_INFO);

EXPECT_FALSE(result.empty());

} else {
EXPECT_EQ(status, SQL_ERROR);

SQLCHAR sql_state[6] = {};
SQLINTEGER native_error = 0;
SQLCHAR message[512] = {};
SQLSMALLINT msg_len = 0;

EXPECT_EQ(
SQLGetDiagRec(SQL_HANDLE_STMT, conn->hstmt, 1, sql_state,
&native_error, message, sizeof(message), &msg_len),
SQL_SUCCESS);

EXPECT_STREQ(reinterpret_cast<char*>(sql_state), "42000");

std::string msg(reinterpret_cast<char*>(message));
EXPECT_NE(msg.find("bigquery.jobs.create permission"), std::string::npos);

// SQLFetch should fail with Function Sequence Error.
status = SQLFetch(conn->hstmt);
EXPECT_EQ(status, SQL_ERROR);

EXPECT_EQ(
SQLGetDiagRec(SQL_HANDLE_STMT, conn->hstmt, 1, sql_state,
&native_error, message, sizeof(message), &msg_len),
SQL_SUCCESS);

msg.assign(reinterpret_cast<char*>(message));
EXPECT_NE(msg.find("Function sequence error"), std::string::npos);

// SQLGetData should also fail with Function Sequence Error.
SQLWCHAR buf[256] = {};
SQLLEN indicator = 0;

status =
SQLGetData(conn->hstmt, 1, SQL_C_WCHAR, buf, sizeof(buf), &indicator);

EXPECT_EQ(status, SQL_ERROR);

EXPECT_EQ(
SQLGetDiagRec(SQL_HANDLE_STMT, conn->hstmt, 1, sql_state,
&native_error, message, sizeof(message), &msg_len),
SQL_SUCCESS);

msg.assign(reinterpret_cast<char*>(message));
EXPECT_NE(msg.find("Function sequence error"), std::string::npos);
}

EXPECT_EQ(Disconnect(conn), SQL_SUCCESS);
}
}

class IgnoreTransactionsTransactionTest
: public ::testing::TestWithParam<
std::tuple<std::string, // IgnoreTransactions
Expand Down Expand Up @@ -4397,6 +4512,7 @@ TEST_P(IgnoreTransactionsTransactionTest, VerifyTransactionBehavior) {
"INSERT INTO " + table_name + " VALUES (1, 'sampledata')";

ASSERT_EQ(SQLExecDirect(conn->hstmt, (SQLCHAR*)insert_query.c_str(), SQL_NTS),

SQL_SUCCESS);

ASSERT_EQ(SQLEndTran(SQL_HANDLE_DBC, conn->hdbc, completion_type),
Expand Down
Loading