feat(BigQuery): Enable session in load job (#6576)

googleapis · Sep 7, 2023 · c27f058 · c27f058
1 parent ad50f74
commit c27f058
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 3 deletions.
diff --git a/BigQuery/src/LoadJobConfiguration.php b/BigQuery/src/LoadJobConfiguration.php
@@ -637,4 +637,46 @@ public function hivePartitioningOptions(array $hivePartitioningOptions)
 
         return $this;
     }
+
+    /**
+     * Decide whether to create a new session with a random id.
+     * The created session id is returned as part of the SessionInfo
+     * within the query statistics.
+     *
+     * Example:
+     * ```
+     * $loadJobConfig->createSession(true);
+     * ```
+     *
+     * @param bool $createSession
+     * @return LoadJobConfiguration
+     */
+    public function createSession(bool $createSession)
+    {
+        $this->config['configuration']['load']['createSession'] = $createSession;
+
+        return $this;
+    }
+
+    /**
+     * Sets connection properties for the load job.
+     *
+     * Example:
+     * ```
+     * $loadJobConfig->connectionProperties([
+     *     'key' => 'session_id',
+     *     'value' => 'sessionId'
+     * ]);
+     * ```
+     *
+     * @param array $connectionProperties
+     * @return LoadJobConfiguration
+     */
+    public function connectionProperties(array $connectionProperties)
+    {
+        $this->config['configuration']['load']['connectionProperties'] =
+            $connectionProperties;
+
+        return $this;
+    }
 }
diff --git a/BigQuery/tests/Snippet/LoadJobConfigurationTest.php b/BigQuery/tests/Snippet/LoadJobConfigurationTest.php
@@ -211,7 +211,18 @@ public function setterDataProvider()
                     'sourceUriPrefix' => 'gs://bucket/path_to_table',
                     'requirePartitionFilter' => false,
                 ]
-            ]
+            ],
+            [
+                'createSession',
+                true
+            ],
+            [
+                'connectionProperties',
+                [
+                    'key' => 'session_id',
+                    'value' => 'sessionId'
+                ]
+            ],
         ];
     }
 }
diff --git a/BigQuery/tests/System/LoadDataAndQueryTest.php b/BigQuery/tests/System/LoadDataAndQueryTest.php
@@ -884,6 +884,47 @@ public function testInsertValidJsonToTable($json)
         $this->assertTrue($insertResponse->isSuccessful());
     }
 
+    public function testLoadWithSession()
+    {
+        $dataset = self::$client->dataset('_SESSION');
+        $table = $dataset->table(uniqid(self::TESTING_PREFIX));
+        $source = __DIR__ . '/data/table-data.json';
+        $schema = [
+            'fields' => [
+                [
+                    'name' => 'name',
+                    'type' => 'STRING'
+                ],
+                [
+                    'name' => 'post_abbr',
+                    'type' => 'STRING'
+                ]
+            ]
+        ];
+        $loadJobConfig = $table->load(fopen($source, 'r'))
+            ->sourceFormat('NEWLINE_DELIMITED_JSON')
+            ->autodetect(true)
+            ->createSession(true);
+
+        $job = self::$client->runJob($loadJobConfig);
+
+        $this->assertTrue(isset($job->info()['statistics']['sessionInfo']['sessionId']));
+
+        $sessionId = $job->info()['statistics']['sessionInfo']['sessionId'];
+        $loadJobConfig = $table->load(fopen($source, 'r'))
+            ->sourceFormat('NEWLINE_DELIMITED_JSON')
+            ->connectionProperties(
+                ['key' => 'session_id', 'value' =>$sessionId]
+            );
+        $job = self::$client->runJob($loadJobConfig);
+        $this->assertFalse(isset($job->info()['status']['errorResult']));
+        $queryJobConfig = self::$client->query(
+            sprintf("CALL BQ.ABORT_SESSION('%s')", $sessionId)
+        );
+        $job = self::$client->runJob($queryJobConfig);
+        $this->assertFalse(isset($job->info()['status']['errorResult']));
+    }
+
     public function invalidJsonProvider()
     {
         return[

diff --git a/BigQuery/tests/Unit/BigQueryClientTest.php b/BigQuery/tests/Unit/BigQueryClientTest.php
@@ -529,6 +529,19 @@ public function jobConfigDataProvider()
                 $expected + ['data' => 'abc'],
                 'insertJobUpload',
                 $uploader->reveal()
+            ],
+            [
+                $expected + ['configuration' => [
+                    'load' => [
+                        'createSession' => true,
+                        'connectionProperties' => [
+                            'key' => 'session_id',
+                            'value' => 'sessionId'
+                        ]
+                    ]
+                ]],
+                'insertJob',
+                $this->jobResponse
             ]
         ];
     }

diff --git a/BigQuery/tests/Unit/LoadJobConfigurationTest.php b/BigQuery/tests/Unit/LoadJobConfigurationTest.php
@@ -103,7 +103,12 @@ public function testFluentSetters()
                 ]
             ],
             'writeDisposition' => 'WRITE_TRUNCATE',
-            'useAvroLogicalTypes' => true
+            'useAvroLogicalTypes' => true,
+            'createSession' => true,
+            'connectionProperties' => [
+                'key' => 'session_id',
+                'value' => 'sessionId'
+            ]
         ];
         $this->expectedConfig['configuration']['load'] = $load
             + $this->expectedConfig['configuration']['load'];
@@ -132,7 +137,12 @@ public function testFluentSetters()
             ->timePartitioning($load['timePartitioning'])
             ->rangePartitioning($load['rangePartitioning'])
             ->writeDisposition($load['writeDisposition'])
-            ->useAvroLogicalTypes($load['useAvroLogicalTypes']);
+            ->useAvroLogicalTypes($load['useAvroLogicalTypes'])
+            ->createSession(true)
+            ->connectionProperties([
+                'key' => 'session_id',
+                'value' =>'sessionId'
+            ]);
 
         $this->assertInstanceOf(LoadJobConfiguration::class, $config);
         $this->assertEquals(