Skip to content

Commit 877b9b5

Browse files
authored
Merge pull request #255 from keboola/vb/DMD-452/BQ-PK-support
DMD-452 - BigQuery add PK support in table utils
2 parents 4557ee4 + b80e693 commit 877b9b5

4 files changed

Lines changed: 262 additions & 3 deletions

File tree

src/Table/Bigquery/BigqueryTableQueryBuilder.php

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ class BigqueryTableQueryBuilder implements TableQueryBuilderInterface
2020
private const INVALID_DEFAULT_VALUE_FOR_NUMERIC_COLUMN = 'invalidDefaultValueForNumericColumn';
2121
private const INVALID_DEFAULT_VALUE_FOR_BOOLAN_COLUMN = 'invalidDefaultValueForBooleanColumn';
2222
private const INVALID_KEY_TO_UPDATE = 'invalidKeyToUpdate';
23+
private const INVALID_PK_FOR_TABLE = 'invalidPKForTable';
2324

2425
/** @var string[] */
2526
protected static array $boolishValues = ['true', 'false', '0', '1'];
@@ -55,11 +56,12 @@ public function getCreateTableCommand(
5556
ColumnCollection $columns,
5657
array $primaryKeys = [],
5758
): string {
58-
assert(count($primaryKeys) === 0, 'primary keys aren\'t supported in BQ');
5959
$columnsSqlDefinitions = [];
60+
$columnNames = [];
6061
/** @var BigqueryColumn $column */
6162
foreach ($columns->getIterator() as $column) {
6263
$columnName = $column->getColumnName();
64+
$columnNames[] = $columnName;
6365
$columnDefinition = $column->getColumnDefinition();
6466

6567
$columnsSqlDefinitions[] = sprintf(
@@ -68,6 +70,30 @@ public function getCreateTableCommand(
6870
$columnDefinition->getSQLDefinition(),
6971
);
7072
}
73+
74+
// check that all PKs are valid columns
75+
$pksNotPresentInColumns = array_diff($primaryKeys, $columnNames);
76+
if ($pksNotPresentInColumns !== []) {
77+
throw new QueryBuilderException(
78+
sprintf(
79+
'Trying to set "%s" as PKs but not present in columns',
80+
implode(',', $pksNotPresentInColumns),
81+
),
82+
self::INVALID_PK_FOR_TABLE,
83+
);
84+
}
85+
86+
if ($primaryKeys !== []) {
87+
$columnsSqlDefinitions[] =
88+
sprintf(
89+
'PRIMARY KEY (%s) NOT ENFORCED',
90+
implode(',', array_map(
91+
static fn($item) => BigqueryQuote::quoteSingleIdentifier($item),
92+
$primaryKeys,
93+
)),
94+
);
95+
}
96+
7197
$columnsSql = implode(",\n", $columnsSqlDefinitions);
7298
return sprintf(
7399
'CREATE TABLE %s.%s

tests/Functional/Bigquery/Schema/BigquerySchemaReflectionTest.php

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,5 +140,81 @@ public function createTableTestFromDefinitionSqlProvider(): Generator
140140
,
141141
'createPrimaryKeys' => true,
142142
];
143+
144+
yield 'single primary key' => [
145+
'definition' => new BigqueryTableDefinition(
146+
$testDb,
147+
$tableName,
148+
false,
149+
new ColumnCollection(
150+
[
151+
BigqueryColumn::createGenericColumn('id'),
152+
BigqueryColumn::createGenericColumn('name'),
153+
],
154+
),
155+
['id'],
156+
),
157+
'query' => <<<EOT
158+
CREATE TABLE `$testDb`.`$tableName`
159+
(
160+
`id` STRING DEFAULT '' NOT NULL,
161+
`name` STRING DEFAULT '' NOT NULL,
162+
PRIMARY KEY (`id`) NOT ENFORCED
163+
);
164+
EOT
165+
,
166+
'createPrimaryKeys' => true,
167+
];
168+
169+
yield 'composite primary key' => [
170+
'definition' => new BigqueryTableDefinition(
171+
$testDb,
172+
$tableName,
173+
false,
174+
new ColumnCollection(
175+
[
176+
BigqueryColumn::createGenericColumn('id'),
177+
BigqueryColumn::createGenericColumn('type'),
178+
BigqueryColumn::createGenericColumn('name'),
179+
],
180+
),
181+
['id', 'type'],
182+
),
183+
'query' => <<<EOT
184+
CREATE TABLE `$testDb`.`$tableName`
185+
(
186+
`id` STRING DEFAULT '' NOT NULL,
187+
`type` STRING DEFAULT '' NOT NULL,
188+
`name` STRING DEFAULT '' NOT NULL,
189+
PRIMARY KEY (`id`,`type`) NOT ENFORCED
190+
);
191+
EOT
192+
,
193+
'createPrimaryKeys' => true,
194+
];
195+
196+
yield 'primary keys not created when flag is false' => [
197+
'definition' => new BigqueryTableDefinition(
198+
$testDb,
199+
$tableName,
200+
false,
201+
new ColumnCollection(
202+
[
203+
BigqueryColumn::createGenericColumn('id'),
204+
BigqueryColumn::createGenericColumn('name'),
205+
],
206+
),
207+
['id'],
208+
),
209+
'query' => <<<EOT
210+
CREATE TABLE `$testDb`.`$tableName`
211+
(
212+
`id` STRING DEFAULT '' NOT NULL,
213+
`name` STRING DEFAULT '' NOT NULL
214+
);
215+
EOT
216+
,
217+
'createPrimaryKeys' => false,
218+
];
143219
}
144220
}

tests/Functional/Bigquery/Table/BigqueryTableQueryBuilderTest.php

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
use Keboola\Datatype\Definition\Bigquery;
1111
use Keboola\TableBackendUtils\Column\Bigquery\BigqueryColumn;
1212
use Keboola\TableBackendUtils\Column\ColumnCollection;
13+
use Keboola\TableBackendUtils\QueryBuilderException;
1314
use Keboola\TableBackendUtils\Table\Bigquery\BigqueryTableQueryBuilder;
1415
use Keboola\TableBackendUtils\Table\Bigquery\BigqueryTableReflection;
1516
use Keboola\TableBackendUtils\TableNotExistsReflectionException;
@@ -48,7 +49,7 @@ public function testGetCreateCommand(
4849
self::TEST_SCHEMA,
4950
self::TABLE_GENERIC,
5051
new ColumnCollection($columns),
51-
[], // primary keys aren't supported in BQ
52+
$primaryKeys,
5253
);
5354

5455
self::assertSame($expectedSql, $sql);
@@ -87,6 +88,67 @@ public function createTableTestSqlProvider(): Generator
8788
EOT
8889
,
8990
];
91+
92+
yield 'single primary key' => [
93+
'cols' => [
94+
BigqueryColumn::createGenericColumn('id'),
95+
BigqueryColumn::createGenericColumn('name'),
96+
],
97+
'primaryKeys' => ['id'],
98+
'expectedColumnNames' => ['id', 'name'],
99+
'expectedPrimaryKeys' => ['id'],
100+
'query' => <<<EOT
101+
CREATE TABLE `$testDb`.`$tableName`
102+
(
103+
`id` STRING DEFAULT '' NOT NULL,
104+
`name` STRING DEFAULT '' NOT NULL,
105+
PRIMARY KEY (`id`) NOT ENFORCED
106+
);
107+
EOT
108+
,
109+
];
110+
111+
yield 'composite primary key' => [
112+
'cols' => [
113+
BigqueryColumn::createGenericColumn('id'),
114+
BigqueryColumn::createGenericColumn('type'),
115+
BigqueryColumn::createGenericColumn('name'),
116+
],
117+
'primaryKeys' => ['id', 'type'],
118+
'expectedColumnNames' => ['id', 'type', 'name'],
119+
'expectedPrimaryKeys' => ['id', 'type'],
120+
'query' => <<<EOT
121+
CREATE TABLE `$testDb`.`$tableName`
122+
(
123+
`id` STRING DEFAULT '' NOT NULL,
124+
`type` STRING DEFAULT '' NOT NULL,
125+
`name` STRING DEFAULT '' NOT NULL,
126+
PRIMARY KEY (`id`,`type`) NOT ENFORCED
127+
);
128+
EOT
129+
,
130+
];
131+
}
132+
133+
public function testInvalidPrimaryKeyThrowsException(): void
134+
{
135+
$this->cleanDataset(self::TEST_SCHEMA);
136+
$this->createDataset(self::TEST_SCHEMA);
137+
138+
$columns = [
139+
BigqueryColumn::createGenericColumn('col1'),
140+
BigqueryColumn::createGenericColumn('col2'),
141+
];
142+
143+
$this->expectException(QueryBuilderException::class);
144+
$this->expectExceptionMessage('Trying to set "nonexistent" as PKs but not present in columns');
145+
146+
$this->qb->getCreateTableCommand(
147+
self::TEST_SCHEMA,
148+
self::TABLE_GENERIC,
149+
new ColumnCollection($columns),
150+
['nonexistent'],
151+
);
90152
}
91153

92154
public function testGetDropTableCommand(): void
@@ -120,7 +182,7 @@ public function testAddAndDropColumn(): void
120182
self::TEST_SCHEMA,
121183
self::TABLE_GENERIC,
122184
new ColumnCollection($columns),
123-
[], // primary keys aren't supported in BQ
185+
[],
124186
);
125187

126188
$this->bqClient->runQuery($this->bqClient->query($sql));

tests/Unit/Table/BigQuery/BigqueryTableQueryBuilderTest.php

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
use Generator;
88
use Keboola\Datatype\Definition\Bigquery;
99
use Keboola\Datatype\Definition\Common;
10+
use Keboola\TableBackendUtils\Column\Bigquery\BigqueryColumn;
11+
use Keboola\TableBackendUtils\Column\ColumnCollection;
1012
use Keboola\TableBackendUtils\QueryBuilderException;
1113
use Keboola\TableBackendUtils\Table\Bigquery\BigqueryTableQueryBuilder;
1214
use Tests\Keboola\TableBackendUtils\Functional\Bigquery\BigqueryBaseCase;
@@ -107,4 +109,97 @@ public function alterColumnCommandInvalidProvider(): Generator
107109
),
108110
];
109111
}
112+
113+
public function testCreateTableWithPrimaryKey(): void
114+
{
115+
$columns = new ColumnCollection([
116+
new BigqueryColumn(
117+
'id',
118+
new Bigquery(Bigquery::TYPE_INTEGER),
119+
),
120+
new BigqueryColumn(
121+
'name',
122+
new Bigquery(Bigquery::TYPE_STRING),
123+
),
124+
]);
125+
126+
$sql = $this->qb->getCreateTableCommand(
127+
'mydataset',
128+
'mytable',
129+
$columns,
130+
['id'],
131+
);
132+
133+
$expectedSql = <<<SQL
134+
CREATE TABLE `mydataset`.`mytable`
135+
(
136+
`id` INTEGER,
137+
`name` STRING,
138+
PRIMARY KEY (`id`) NOT ENFORCED
139+
);
140+
SQL;
141+
142+
$this->assertSame($expectedSql, $sql);
143+
}
144+
145+
public function testCreateTableWithCompositePrimaryKey(): void
146+
{
147+
$columns = new ColumnCollection([
148+
new BigqueryColumn(
149+
'id',
150+
new Bigquery(Bigquery::TYPE_INTEGER),
151+
),
152+
new BigqueryColumn(
153+
'type',
154+
new Bigquery(Bigquery::TYPE_STRING),
155+
),
156+
new BigqueryColumn(
157+
'name',
158+
new Bigquery(Bigquery::TYPE_STRING),
159+
),
160+
]);
161+
162+
$sql = $this->qb->getCreateTableCommand(
163+
'mydataset',
164+
'mytable',
165+
$columns,
166+
['id', 'type'],
167+
);
168+
169+
$expectedSql = <<<SQL
170+
CREATE TABLE `mydataset`.`mytable`
171+
(
172+
`id` INTEGER,
173+
`type` STRING,
174+
`name` STRING,
175+
PRIMARY KEY (`id`,`type`) NOT ENFORCED
176+
);
177+
SQL;
178+
179+
$this->assertSame($expectedSql, $sql);
180+
}
181+
182+
public function testCreateTableWithInvalidPrimaryKey(): void
183+
{
184+
$columns = new ColumnCollection([
185+
new BigqueryColumn(
186+
'id',
187+
new Bigquery(Bigquery::TYPE_INTEGER),
188+
),
189+
new BigqueryColumn(
190+
'name',
191+
new Bigquery(Bigquery::TYPE_STRING),
192+
),
193+
]);
194+
195+
$this->expectException(QueryBuilderException::class);
196+
$this->expectExceptionMessage('Trying to set "nonexistent" as PKs but not present in columns');
197+
198+
$this->qb->getCreateTableCommand(
199+
'mydataset',
200+
'mytable',
201+
$columns,
202+
['nonexistent'],
203+
);
204+
}
110205
}

0 commit comments

Comments
 (0)