packages/ztd-query-postgres/tests/Unit/Rewrite/Sampling/PgSqlTableSampleRewriterTest.php

1<?php
2
3declare(strict_types=1);
4
5namespace Tests\Unit\Rewrite\Sampling;
6
7use PHPUnit\Framework\Attributes\CoversClass;
8use PHPUnit\Framework\Attributes\UsesClass;
9use PHPUnit\Framework\TestCase;
10use ZtdQuery\Exception\UnsupportedSqlException;
11use ZtdQuery\Platform\Postgres\Rewrite\Sampling\PgSqlTableSampleRewriter;
12use ZtdQuery\Platform\Postgres\Sql\PgSqlIdentifierQuoter;
13use ZtdQuery\Platform\Postgres\Sql\Sampling\PgSqlTableSample;
14use ZtdQuery\Platform\Postgres\Sql\Sampling\PgSqlTableSampleParser;
15
16#[CoversClass(PgSqlTableSampleRewriter::class)]
17#[UsesClass(PgSqlTableSampleParser::class)]
18#[UsesClass(PgSqlTableSample::class)]
19#[UsesClass(PgSqlIdentifierQuoter::class)]
20#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Relation\PgSqlSelectRelationParser::class)]
21#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\PgSqlLexerProfile::class)]
22#[CoversClass(\ZtdQuery\Platform\Postgres\Rewrite\Sampling\SampleProjection::class)]
23#[CoversClass(\ZtdQuery\Platform\Postgres\Rewrite\Sampling\TableColumns::class)]
24#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Relation\FromClause::class)]
25#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Relation\RelationReference::class)]
26#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Sampling\SampleClause::class)]
27#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Sampling\SampleTokens::class)]
28#[UsesClass(\ZtdQuery\Platform\Postgres\Sql\Sampling\PgSqlTableSampleMethod::class)]
29final class PgSqlTableSampleRewriterTest extends TestCase
30{
31    public function testRewritesBernoulliSampleAsDerivedShadowRelation(): void
32    {
33        $result = (new PgSqlTableSampleRewriter())->rewrite(
34            'SELECT sampled.id FROM data AS sampled TABLESAMPLE BERNOULLI ($1)',
35            ['data' => ['columns' => ['id', 'value']]],
36        );
37
38        self::assertStringNotContainsString('TABLESAMPLE', $result);
39        self::assertStringContainsString('ROW_NUMBER() OVER () AS "__ztd_sample_ordinal"', $result);
40        self::assertStringContainsString('FROM data', $result);
41        self::assertStringContainsString('CAST(($1) AS DOUBLE PRECISION)', $result);
42        self::assertStringContainsString('random() AS "__ztd_sample_seed"', $result);
43        self::assertStringEndsWith(' AS sampled', $result);
44        self::assertSame(
45            '924604101f2f5928abaa874fa6d54f176e4e6d6afa2a81334267eea904a147eb',
46            hash('sha256', $result),
47        );
48    }
49
50    public function testRepeatableUsesProvidedSeedAndSystemUsesOneVirtualBlock(): void
51    {
52        $result = (new PgSqlTableSampleRewriter())->rewrite(
53            'SELECT * FROM data TABLESAMPLE SYSTEM (50) REPEATABLE (7.5)',
54            ['data' => ['columns' => ['id']]],
55        );
56
57        self::assertStringContainsString('CAST((7.5) AS DOUBLE PRECISION)', $result);
58        self::assertStringContainsString('MD5(CAST(0 AS TEXT)', $result);
59        self::assertStringEndsWith(' AS "data"', $result);
60        self::assertSame(
61            '0b97b0d5fe25795951b26924624bea41df8218e2f895764d11009aee7da164d6',
62            hash('sha256', $result),
63        );
64    }
65
66    public function testMultipleSamplesAreReplacedWithoutOffsetCorruption(): void
67    {
68        $result = (new PgSqlTableSampleRewriter())->rewrite(
69            'SELECT * FROM data TABLESAMPLE BERNOULLI (100), logs TABLESAMPLE SYSTEM (0)',
70            [
71                'data' => ['columns' => ['id']],
72                'logs' => ['columns' => ['id', 'message']],
73            ],
74        );
75
76        self::assertStringNotContainsString('TABLESAMPLE', $result);
77        self::assertSame(2, substr_count($result, 'ROW_NUMBER() OVER ()'));
78        self::assertStringContainsString('FROM data', $result);
79        self::assertStringContainsString('FROM logs', $result);
80        self::assertSame(
81            'd0e77ae3bc03c3a002cd20bbd7de31a6cce7d73b2ae8cc1e3764f111b931c4e7',
82            hash('sha256', $result),
83        );
84    }
85
86    public function testOrdinarySelectIsUnchanged(): void
87    {
88        $sql = 'SELECT * FROM data';
89
90        self::assertSame($sql, (new PgSqlTableSampleRewriter())->rewrite($sql, []));
91    }
92
93    public function testRejectsSampleWhenColumnsAreUnknown(): void
94    {
95        $this->expectException(UnsupportedSqlException::class);
96        $this->expectExceptionMessage('Cannot determine columns');
97
98        (new PgSqlTableSampleRewriter())->rewrite(
99            'SELECT * FROM data TABLESAMPLE BERNOULLI (50)',
100            ['data' => ['columns' => []]],
101        );
102    }
103
104    public function testIgnoresNonStringColumnMetadataAndReindexesNames(): void
105    {
106        $result = (new PgSqlTableSampleRewriter())->rewrite(
107            'SELECT * FROM data TABLESAMPLE BERNOULLI (50)',
108            ['data' => ['columns' => ['id', 123, 'value']]],
109        );
110
111        self::assertStringContainsString('SELECT "id", "value"', $result);
112        self::assertStringNotContainsString('123', $result);
113    }
114}
115