use Retry function for BQ from utils

keboola · Feb 5, 2025 · 2a55a0f · 2a55a0f
1 parent 82e1aec
commit 2a55a0f
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 17 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()),
             '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,6 +14,7 @@
 use Keboola\TableBackendUtils\Escaping\Exasol\ExasolQuote;
 use Keboola\TableBackendUtils\Escaping\Teradata\TeradataQuote;
 use Keboola\Test\StorageApiTestCase;
+use Psr\Log\NullLogger;
 
 trait WorkspaceConnectionTrait
 {
@@ -146,9 +150,13 @@ private function getDbConnectionTeradata(array $connection): DBALConnection
 
     private function getDbConnectionBigquery(array $connection): 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()),
+            'requestTimeout' => 120,
+            'retries' => 10,
+        ], 'sapiTest');
 
         $bqClient->runQuery(
             $bqClient->query('SELECT SESSION_USER() AS USER'),

diff --git a/tests/Backend/Workspaces/Backend/BigqueryWorkspaceBackend.php b/tests/Backend/Workspaces/Backend/BigqueryWorkspaceBackend.php
@@ -6,38 +6,38 @@
 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;
 
     /**
      * @param array $workspace
      */
-    public function __construct(array $workspace)
+    public function __construct(array $workspace, int $retriesCount)
     {
         $bqClient = new BigQueryClientWrapper([
             'keyFile' => $workspace['connection']['credentials'],
+            'restRetryFunction' => Retry::getRestRetryFunction(new NullLogger()),
             'requestTimeout' => 120,
-            'retries' => 20,
-        ]);
+            'retries' => $retriesCount,
+        ], 'sapiTest');
 
-        assert($bqClient instanceof BigQueryClient);
         $this->bqClient = $bqClient;
         $this->schema = $workspace['connection']['schema'];
     }

diff --git a/tests/Backend/Workspaces/Backend/WorkspaceBackendFactory.php b/tests/Backend/Workspaces/Backend/WorkspaceBackendFactory.php
@@ -16,7 +16,7 @@ class WorkspaceBackendFactory
     /**
      * @param array<mixed> $workspace
      */
-    public static function createWorkspaceBackend(array $workspace, bool $useDBAL = false): WorkspaceBackend
+    public static function createWorkspaceBackend(array $workspace, bool $useDBAL = false, int $retriesCount = 10): WorkspaceBackend
     {
         switch ($workspace['connection']['backend']) {
             case StorageApiTestCase::BACKEND_REDSHIFT:
@@ -33,7 +33,7 @@ public static function createWorkspaceBackend(array $workspace, bool $useDBAL =
             case StorageApiTestCase::BACKEND_TERADATA:
                 return new TeradataWorkspaceBackend($workspace);
             case StorageApiTestCase::BACKEND_BIGQUERY:
-                return new BigqueryWorkspaceBackend($workspace);
+                return new BigqueryWorkspaceBackend($workspace, $retriesCount);
             default:
                 throw new Exception($workspace['connection']['backend'] . ' workspaces are not supported.');
         }

diff --git a/tests/Backend/Workspaces/WorkspacesTest.php b/tests/Backend/Workspaces/WorkspacesTest.php
@@ -95,7 +95,7 @@ public function testWorkspaceCreate(bool $async): void
         $tokenInfo = $this->_client->verifyToken();
         $this->assertEquals($tokenInfo['owner']['defaultBackend'], $connection['backend']);
 
-        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace);
+        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace, false, 2);
 
         if ($workspaceWithSnowflakeBackend) {
             $db = $backend->getDb();
@@ -254,7 +254,7 @@ public function testWorkspacePasswordReset(): void
         $tokenInfo = $this->_client->verifyToken();
         $this->assertEquals($tokenInfo['owner']['defaultBackend'], $connection['backend']);
 
-        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace);
+        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace, false, 2);
         $backend->dropTableIfExists('mytable');
 
         $backend->createTable('mytable', ['amount' => $this->getColumnAmountType($connection['backend'])]);

diff --git a/tests/Common/BranchWorkspacesTest.php b/tests/Common/BranchWorkspacesTest.php
@@ -63,7 +63,7 @@ public function testWorkspaceCreate(bool $async): void
         $tokenInfo = $this->_client->verifyToken();
         $this->assertEquals($tokenInfo['owner']['defaultBackend'], $connection['backend']);
 
-        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace);
+        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace, false, 2);
         $backend->createTable('mytable', ['amount' => ($connection['backend'] === self::BACKEND_SNOWFLAKE) ? 'NUMBER' : 'VARCHAR']);
 
         $tableNames = $backend->getTables();
@@ -124,7 +124,7 @@ public function testWorkspacePasswordReset(): void
         $tokenInfo = $this->_client->verifyToken();
         $this->assertEquals($tokenInfo['owner']['defaultBackend'], $connection['backend']);
 
-        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace);
+        $backend = WorkspaceBackendFactory::createWorkspaceBackend($workspace, false, 2);
         $backend->dropTableIfExists('mytable');
         $backend->createTable('mytable', ['amount' => ($connection['backend'] === self::BACKEND_SNOWFLAKE) ? 'NUMBER' : 'VARCHAR']);