use Retry function for BQ from utils

keboola · Feb 7, 2025 · a602a63 · a602a63
1 parent e52f433
commit a602a63
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 13 deletions.
diff --git a/tests/Backend/ExternalBuckets/BigqueryRegisterBucketTest.php b/tests/Backend/ExternalBuckets/BigqueryRegisterBucketTest.php
@@ -18,13 +18,15 @@
 use Keboola\StorageApi\TableExporter;
 use Keboola\StorageApi\Workspaces;
 use Keboola\TableBackendUtils\Connection\Bigquery\BigQueryClientWrapper;
+use Keboola\TableBackendUtils\Connection\Bigquery\Retry;
 use Keboola\TableBackendUtils\Escaping\Bigquery\BigqueryQuote;
 use Keboola\Test\Backend\Workspaces\Backend\WorkspaceBackend;
 use Keboola\Test\Backend\Workspaces\Backend\WorkspaceBackendFactory;
 use Keboola\Test\ClientProvider\ClientProvider;
 use Keboola\Test\ClientProvider\TestSetupHelper;
 use Keboola\Test\Utils\EventsQueryBuilder;
 use LogicException;
+use Psr\Log\NullLogger;
 
 // Tests for registering an external bucket that contains different types of tables.
 // Available table types https://cloud.google.com/bigquery/docs/information-schema-tables#schema:
@@ -1528,10 +1530,11 @@ public function getBigQueryClient(array $externalCredentials): BigQueryClient
     {
         return new BigQueryClientWrapper([
             'keyFile' => $externalCredentials,
+            'restRetryFunction' => Retry::getRestRetryFunction(new NullLogger(), true),
             'requestTimeout' => 120,
-            'retries' => 20,
+            'retries' => 10,
             'location' => $this->region,
-        ]);
+        ], 'sapitest');
     }
 
     private function createExternalTable(WorkspaceBackend $db, string $schemaName): void

diff --git a/tests/Backend/WorkspaceConnectionTrait.php b/tests/Backend/WorkspaceConnectionTrait.php
@@ -3,6 +3,9 @@
 namespace Keboola\Test\Backend;
 
 use Google\Cloud\BigQuery\BigQueryClient;
+use Keboola\StorageDriver\BigQuery\CredentialsHelper;
+use Keboola\TableBackendUtils\Connection\Bigquery\BigQueryClientWrapper;
+use Keboola\TableBackendUtils\Connection\Bigquery\Retry;
 use Keboola\TableBackendUtils\Connection\Snowflake\SnowflakeConnectionFactory;
 use Keboola\TableBackendUtils\Connection\Teradata\TeradataConnection;
 use Doctrine\DBAL\Connection as DBALConnection;
@@ -11,13 +14,14 @@
 use Keboola\TableBackendUtils\Escaping\Exasol\ExasolQuote;
 use Keboola\TableBackendUtils\Escaping\Teradata\TeradataQuote;
 use Keboola\Test\StorageApiTestCase;
+use Psr\Log\NullLogger;
 
 trait WorkspaceConnectionTrait
 {
     /**
      * @return SnowflakeConnection|DBALConnection|\PDO|BigQueryClient
      */
-    private function getDbConnection(array $connection)
+    private function getDbConnection(array $connection, bool $bigQueryShouldRetry401 = true)
     {
         switch ($connection['backend']) {
             case StorageApiTestCase::BACKEND_SNOWFLAKE:
@@ -31,7 +35,7 @@ private function getDbConnection(array $connection)
             case StorageApiTestCase::BACKEND_TERADATA:
                 return $this->getDbConnectionTeradata($connection);
             case StorageApiTestCase::BACKEND_BIGQUERY:
-                return $this->getDbConnectionBigquery($connection);
+                return $this->getDbConnectionBigquery($connection, $bigQueryShouldRetry401);
         }
 
         throw new \Exception('Unsupported Backend for workspaces');
@@ -144,11 +148,15 @@ private function getDbConnectionTeradata(array $connection): DBALConnection
         return $db;
     }
 
-    private function getDbConnectionBigquery(array $connection): BigQueryClient
+    private function getDbConnectionBigquery(array $connection, bool $bigQueryShouldRetry401 = true): BigQueryClient
     {
-        $bqClient = new BigQueryClient([
+        // note: the close method is not used in this client
+        $bqClient = new BigQueryClientWrapper([
             'keyFile' => $connection['credentials'],
-        ]);
+            'restRetryFunction' => Retry::getRestRetryFunction(new NullLogger(), $bigQueryShouldRetry401),
+            'requestTimeout' => 120,
+            'retries' => 20,
+        ], 'sapitest');
 
         $bqClient->runQuery(
             $bqClient->query('SELECT SESSION_USER() AS USER'),

diff --git a/tests/Backend/WorkspaceCredentialsAssertTrait.php b/tests/Backend/WorkspaceCredentialsAssertTrait.php
@@ -23,7 +23,7 @@ private function assertCredentialsShouldNotWork($connection): void
         try {
             $proxy = new RetryProxy($retryPolicy, new ExponentialBackOffPolicy());
             $proxy->call(function () use ($connection) {
-                $this->getDbConnection($connection);
+                $this->getDbConnection($connection, false);
                 throw new \Exception(self::$RETRY_FAIL_MESSAGE);
             });
         } catch (\Doctrine\DBAL\Driver\Exception $e) {

diff --git a/tests/Backend/Workspaces/Backend/BigqueryWorkspaceBackend.php b/tests/Backend/Workspaces/Backend/BigqueryWorkspaceBackend.php
@@ -6,23 +6,23 @@
 use Google\Cloud\BigQuery\Dataset;
 use Keboola\Datatype\Definition\Bigquery;
 use Keboola\StorageApi\Exception;
-use Keboola\StorageDriver\BigQuery\Client\BigQuery\Retry;
-use Keboola\StorageDriver\BigQuery\CredentialsHelper;
 use Keboola\TableBackendUtils\Column\Bigquery\BigqueryColumn;
 use Keboola\TableBackendUtils\Column\ColumnCollection;
 use Keboola\TableBackendUtils\Connection\Bigquery\BigQueryClientWrapper;
+use Keboola\TableBackendUtils\Connection\Bigquery\Retry;
 use Keboola\TableBackendUtils\Escaping\Bigquery\BigqueryQuote;
 use Keboola\TableBackendUtils\Schema\Bigquery\BigquerySchemaReflection;
 use Keboola\TableBackendUtils\Table\Bigquery\BigqueryTableQueryBuilder;
 use Keboola\TableBackendUtils\Table\Bigquery\BigqueryTableReflection;
 use Keboola\TableBackendUtils\View\ViewReflectionInterface;
 use PDO;
+use Psr\Log\NullLogger;
 use Retry\Policy\SimpleRetryPolicy;
 use Retry\RetryProxy;
 
 class BigqueryWorkspaceBackend implements WorkspaceBackend
 {
-    private BigQueryClient $bqClient;
+    private BigQueryClientWrapper $bqClient;
 
     private string $schema;
 
@@ -33,11 +33,11 @@ public function __construct(array $workspace)
     {
         $bqClient = new BigQueryClientWrapper([
             'keyFile' => $workspace['connection']['credentials'],
+            'restRetryFunction' => Retry::getRestRetryFunction(new NullLogger(), true),
             'requestTimeout' => 120,
             'retries' => 20,
-        ]);
+        ], 'sapitest');
 
-        assert($bqClient instanceof BigQueryClient);
         $this->bqClient = $bqClient;
         $this->schema = $workspace['connection']['schema'];
     }