Skip to content

Commit 87e5d8c

Browse files
dveritedverite
committed
fix: use proper quoting for table names
When table names contain special characters that require identifiers to be quoted, the code was not quoting properly. In particular: - calling Sprintf("%s.%s") with unquoted table and schema names is not safe to build SQL statements. pgx.Identifier.Sanitize() is used instead. - Casting to regclass to convert OID to names is too complicated because it can produce quoted or unquoted names depending on the context. Use pg_class and pg_namespace instead. - Casting to regclass to convert names to OID depends on the search_path, and also requires pre-quoting, which makes it difficult to use properly with user-supplied arguments. Use pg_class and pg_namespace instead. Signed-off-by: Daniel Vérité <dverite@gmail.com>
1 parent 67925e3 commit 87e5d8c

6 files changed

Lines changed: 115 additions & 44 deletions

File tree

‎internal/infra/postgresql/partition.go‎

Lines changed: 41 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,13 @@ type PartitionResult struct {
2424
}
2525

2626
func (p Postgres) IsPartitionAttached(schema, table string) (exists bool, err error) {
27-
query := `SELECT EXISTS(
28-
SELECT 1 FROM pg_inherits WHERE inhrelid = $1::regclass
29-
)`
27+
query := `SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits WHERE inhrelid = (SELECT c.oid
28+
FROM pg_catalog.pg_class c
29+
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
30+
WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind='r'))
31+
`
3032

31-
err = p.conn.QueryRow(p.ctx, query, fmt.Sprintf("%s.%s", schema, table)).Scan(&exists)
33+
err = p.conn.QueryRow(p.ctx, query, schema, table).Scan(&exists)
3234
if err != nil {
3335
return false, fmt.Errorf("failed to check partition attachment: %w", err)
3436
}
@@ -37,7 +39,10 @@ func (p Postgres) IsPartitionAttached(schema, table string) (exists bool, err er
3739
}
3840

3941
func (p Postgres) AttachPartition(schema, table, parent, lowerBound, upperBound string) error {
40-
query := fmt.Sprintf("ALTER TABLE %s.%s ATTACH PARTITION %s.%s FOR VALUES FROM ('%s') TO ('%s')", schema, parent, schema, table, lowerBound, upperBound)
42+
query := fmt.Sprintf("ALTER TABLE %s ATTACH PARTITION %s FOR VALUES FROM ('%s') TO ('%s')",
43+
pgx.Identifier{schema, parent}.Sanitize(),
44+
pgx.Identifier{schema, table}.Sanitize(),
45+
lowerBound, upperBound)
4146
p.logger.Debug("Attach partition", "query", query, "schema", schema, "table", table)
4247

4348
_, err := p.conn.Exec(p.ctx, query)
@@ -52,7 +57,10 @@ func (p Postgres) AttachPartition(schema, table, parent, lowerBound, upperBound
5257
// The partition still exists as standalone table after detaching
5358
// More info: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DETACH-PARTITION
5459
func (p Postgres) DetachPartitionConcurrently(schema, table, parent string) error {
55-
query := fmt.Sprintf(`ALTER TABLE %s.%s DETACH PARTITION %s.%s CONCURRENTLY`, schema, parent, schema, table)
60+
query := fmt.Sprintf("ALTER TABLE %s DETACH PARTITION %s CONCURRENTLY",
61+
pgx.Identifier{schema, parent}.Sanitize(),
62+
pgx.Identifier{schema, table}.Sanitize())
63+
5664
p.logger.Debug("Detach partition", "schema", schema, "table", table, "query", query, "parent_table", parent)
5765

5866
_, err := p.conn.Exec(p.ctx, query)
@@ -67,7 +75,9 @@ func (p Postgres) DetachPartitionConcurrently(schema, table, parent string) erro
6775
// It's required when a partition is in "detach pending" status.
6876
// More info: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DETACH-PARTITION
6977
func (p Postgres) FinalizePartitionDetach(schema, table, parent string) error {
70-
query := fmt.Sprintf(`ALTER TABLE %s.%s DETACH PARTITION %s.%s FINALIZE`, schema, parent, schema, table)
78+
query := fmt.Sprintf(`ALTER TABLE %s DETACH PARTITION %s FINALIZE`,
79+
pgx.Identifier{schema, parent}.Sanitize(),
80+
pgx.Identifier{schema, table}.Sanitize())
7181
p.logger.Debug("finialize detach partition", "schema", schema, "table", table, "query", query, "parent_table", parent)
7282

7383
_, err := p.conn.Exec(p.ctx, query)
@@ -79,28 +89,33 @@ func (p Postgres) FinalizePartitionDetach(schema, table, parent string) error {
7989
}
8090

8191
func (p Postgres) ListPartitions(schema, table string) (partitions []PartitionResult, err error) {
82-
query := fmt.Sprintf(`
92+
query := `
8393
WITH parts as (
8494
SELECT
85-
relnamespace::regnamespace as schema,
86-
c.oid::pg_catalog.regclass AS part_name,
95+
n.nspname as schema,
96+
c.relname AS part_name,
8797
regexp_match(pg_get_expr(c.relpartbound, c.oid),
8898
'FOR VALUES FROM \(''(.*)''\) TO \(''(.*)''\)') AS bounds
8999
FROM
90100
pg_catalog.pg_class c JOIN pg_catalog.pg_inherits i ON (c.oid = i.inhrelid)
91-
WHERE i.inhparent = '%s.%s'::regclass
101+
JOIN pg_catalog.pg_namespace n ON (c.relnamespace = n.oid)
102+
WHERE i.inhparent = (SELECT c.oid
103+
FROM pg_catalog.pg_class c
104+
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
105+
WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind='p' -- parent
106+
)
92107
AND c.relkind='r'
93108
)
94109
SELECT
95110
schema,
96111
part_name as name,
97-
'%s' as parentTable,
112+
$2 as parentTable,
98113
bounds[1]::text AS lowerBound,
99114
bounds[2]::text AS upperBound
100115
FROM parts
101-
ORDER BY part_name;`, schema, table, table)
116+
ORDER BY part_name`
102117

103-
rows, err := p.conn.Query(p.ctx, query)
118+
rows, err := p.conn.Query(p.ctx, query, schema, table)
104119
if err != nil {
105120
return nil, fmt.Errorf("failed to list partitions: %w", err)
106121
}
@@ -119,12 +134,15 @@ func (p Postgres) GetPartitionSettings(schema, table string) (strategy, key stri
119134
// pg_get_partkeydef() is a system function returning the definition of a partitioning key
120135
// It return a text string: <partitioningStrategy> (<partitioning key definition>)
121136
// Example for RANGE (created_at)
122-
query := fmt.Sprintf(`
137+
query := `
123138
SELECT regexp_match(partkeydef, '^(.*) \((.*)\)$')
124-
FROM pg_catalog.pg_get_partkeydef('%s.%s'::regclass) as partkeydef
125-
`, schema, table)
139+
FROM pg_catalog.pg_get_partkeydef((SELECT c.oid
140+
FROM pg_catalog.pg_class c
141+
JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
142+
WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind='p')) as partkeydef
143+
`
126144

127-
err = p.conn.QueryRow(p.ctx, query).Scan(&partkeydef)
145+
err = p.conn.QueryRow(p.ctx, query, schema, table).Scan(&partkeydef)
128146
if err != nil {
129147
p.logger.Warn("failed to get partitioning key", "error", err, "schema", schema, "table", table)
130148

@@ -153,7 +171,8 @@ func (p Postgres) SetPartitionReplicaIdentity(schema, table, parent string) erro
153171
}
154172

155173
if replIdent == "f" { // replica identity = full
156-
queryFull := fmt.Sprintf("ALTER TABLE %s.%s REPLICA IDENTITY FULL", schema, table)
174+
queryFull := fmt.Sprintf("ALTER TABLE %s REPLICA IDENTITY FULL",
175+
pgx.Identifier{schema, table}.Sanitize())
157176
p.logger.Debug("Set identity full", "query", queryFull)
158177

159178
_, err = p.conn.Exec(p.ctx, queryFull)
@@ -184,7 +203,9 @@ SELECT c_idx_child.relname
184203
return fmt.Errorf("failed to find the child index for the new partition: %w", err)
185204
}
186205

187-
queryAlter := fmt.Sprintf("ALTER TABLE %s.%s REPLICA IDENTITY USING INDEX %s", schema, table, indexName)
206+
queryAlter := fmt.Sprintf("ALTER TABLE %s REPLICA IDENTITY USING INDEX %s",
207+
pgx.Identifier{schema, table}.Sanitize(),
208+
pgx.Identifier{indexName}.Sanitize())
188209
p.logger.Debug("Set replica identity", "schema", schema, "table", table, "index", indexName, "query", queryAlter)
189210

190211
_, err := p.conn.Exec(p.ctx, queryAlter)

‎internal/infra/postgresql/partition_test.go‎

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"testing"
77

8+
"github.com/jackc/pgx/v5"
89
"github.com/pashagolub/pgxmock/v3"
910
"github.com/qonto/postgresql-partition-manager/internal/infra/postgresql"
1011
"github.com/stretchr/testify/assert"
@@ -15,29 +16,29 @@ func generateTable(t *testing.T) (schema, table, fullQualifiedTable, parent stri
1516

1617
schema = "public"
1718
table = "my_table"
18-
fullQualifiedTable = fmt.Sprintf("%s.%s", schema, table)
19+
fullQualifiedTable = pgx.Identifier{schema, table}.Sanitize()
1920
parent = "my_parent_table"
2021

2122
return
2223
}
2324

2425
func TestIsPartitionAttached(t *testing.T) {
25-
schema, table, fullQualifiedTable, _ := generateTable(t)
26+
schema, table, _, _ := generateTable(t)
2627

2728
mock, p := setupMock(t, pgxmock.QueryMatcherRegexp)
2829
query := "SELECT EXISTS"
2930

30-
mock.ExpectQuery(query).WithArgs(fullQualifiedTable).WillReturnRows(mock.NewRows([]string{"EXISTS"}).AddRow(true))
31+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnRows(mock.NewRows([]string{"EXISTS"}).AddRow(true))
3132
exists, err := p.IsPartitionAttached(schema, table)
3233
assert.Nil(t, err, "IsPartitionAttached should succeed")
3334
assert.True(t, exists, "Table should be attached")
3435

35-
mock.ExpectQuery(query).WithArgs(fullQualifiedTable).WillReturnRows(mock.NewRows([]string{"EXISTS"}).AddRow(false))
36+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnRows(mock.NewRows([]string{"EXISTS"}).AddRow(false))
3637
exists, err = p.IsPartitionAttached(schema, table)
3738
assert.Nil(t, err, "IsPartitionAttached should succeed")
3839
assert.False(t, exists, "Table should not be attached")
3940

40-
mock.ExpectQuery(query).WithArgs(fullQualifiedTable).WillReturnError(ErrPostgreSQLConnectionFailure)
41+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnError(ErrPostgreSQLConnectionFailure)
4142
_, err = p.IsPartitionAttached(schema, table)
4243
assert.Error(t, err, "IsPartitionAttached should fail")
4344
}
@@ -48,7 +49,10 @@ func TestAttachPartition(t *testing.T) {
4849
upperBound := "2024-01-31"
4950

5051
mock, p := setupMock(t, pgxmock.QueryMatcherEqual)
51-
query := fmt.Sprintf(`ALTER TABLE %s.%s ATTACH PARTITION %s.%s FOR VALUES FROM ('%s') TO ('%s')`, schema, parent, schema, table, lowerBound, upperBound)
52+
query := fmt.Sprintf(`ALTER TABLE %s ATTACH PARTITION %s FOR VALUES FROM ('%s') TO ('%s')`,
53+
pgx.Identifier{schema, parent}.Sanitize(),
54+
pgx.Identifier{schema, table}.Sanitize(),
55+
lowerBound, upperBound)
5256

5357
mock.ExpectExec(query).WillReturnResult(pgxmock.NewResult("ALTER", 1))
5458
err := p.AttachPartition(schema, table, parent, lowerBound, upperBound)
@@ -63,7 +67,9 @@ func TestDetachPartitionConcurrently(t *testing.T) {
6367
schema, table, _, parent := generateTable(t)
6468

6569
mock, p := setupMock(t, pgxmock.QueryMatcherEqual)
66-
query := fmt.Sprintf(`ALTER TABLE %s.%s DETACH PARTITION %s.%s CONCURRENTLY`, schema, parent, schema, table)
70+
query := fmt.Sprintf(`ALTER TABLE %s DETACH PARTITION %s CONCURRENTLY`,
71+
pgx.Identifier{schema, parent}.Sanitize(),
72+
pgx.Identifier{schema, table}.Sanitize())
6773

6874
mock.ExpectExec(query).WillReturnResult(pgxmock.NewResult("ALTER", 1))
6975
err := p.DetachPartitionConcurrently(schema, table, parent)
@@ -79,7 +85,9 @@ func TestFinalizePartitionDetach(t *testing.T) {
7985

8086
mock, p := setupMock(t, pgxmock.QueryMatcherEqual)
8187

82-
query := fmt.Sprintf(`ALTER TABLE %s.%s DETACH PARTITION %s.%s FINALIZE`, schema, parent, schema, table)
88+
query := fmt.Sprintf(`ALTER TABLE %s DETACH PARTITION %s FINALIZE`,
89+
pgx.Identifier{schema, parent}.Sanitize(),
90+
pgx.Identifier{schema, table}.Sanitize())
8391

8492
mock.ExpectExec(query).WillReturnResult(pgxmock.NewResult("ALTER", 1))
8593
err := p.FinalizePartitionDetach(schema, table, parent)
@@ -99,18 +107,18 @@ func TestGetPartitionSettings(t *testing.T) {
99107

100108
query := `SELECT regexp_match`
101109

102-
mock.ExpectQuery(query).WillReturnRows(mock.NewRows([]string{"partkeydef"}).AddRow([]string{expectedStrategy, expectedKey}))
110+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnRows(mock.NewRows([]string{"partkeydef"}).AddRow([]string{expectedStrategy, expectedKey}))
103111
strategy, key, err := p.GetPartitionSettings(schema, table)
104112
assert.Nil(t, err, "GetPartitionSettings should succeed")
105113
assert.Equal(t, strategy, expectedStrategy, "Strategy should match")
106114
assert.Equal(t, key, expectedKey, "Key should match")
107115

108-
mock.ExpectQuery(query).WillReturnRows(mock.NewRows([]string{"partkeydef"}).AddRow([]string{}))
116+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnRows(mock.NewRows([]string{"partkeydef"}).AddRow([]string{}))
109117
_, _, err = p.GetPartitionSettings(schema, table)
110118
assert.Error(t, err, "GetPartitionSettings should fail")
111119
assert.ErrorIs(t, err, postgresql.ErrTableIsNotPartitioned)
112120

113-
mock.ExpectQuery(query).WillReturnError(ErrPostgreSQLConnectionFailure)
121+
mock.ExpectQuery(query).WithArgs(schema, table).WillReturnError(ErrPostgreSQLConnectionFailure)
114122
_, _, err = p.GetPartitionSettings(schema, table)
115123
assert.Error(t, err, "GetPartitionSettings should fail")
116124
}
@@ -142,17 +150,17 @@ func TestListPartitions(t *testing.T) {
142150
for _, p := range expectedPartitions {
143151
rows.AddRow(p.Schema, p.Name, p.ParentTable, p.LowerBound, p.UpperBound)
144152
}
145-
mock.ExpectQuery(query).WillReturnRows(rows)
146-
result, err := p.ListPartitions(schema, table)
153+
mock.ExpectQuery(query).WithArgs(schema, parent).WillReturnRows(rows)
154+
result, err := p.ListPartitions(schema, parent)
147155
assert.Nil(t, err, "ListPartitions should succeed")
148156
assert.Equal(t, result, expectedPartitions, "Partitions should be match")
149157

150158
rows = mock.NewRows([]string{"invalidColumn"}).AddRow("invalidColumn")
151-
mock.ExpectQuery(query).WillReturnRows(rows)
152-
_, err = p.ListPartitions(schema, table)
159+
mock.ExpectQuery(query).WithArgs(schema, parent).WillReturnRows(rows)
160+
_, err = p.ListPartitions(schema, parent)
153161
assert.Error(t, err, "ListPartitions should fail")
154162

155-
mock.ExpectQuery(query).WillReturnError(ErrPostgreSQLConnectionFailure)
156-
_, err = p.ListPartitions(schema, table)
163+
mock.ExpectQuery(query).WithArgs(schema, parent).WillReturnError(ErrPostgreSQLConnectionFailure)
164+
_, err = p.ListPartitions(schema, parent)
157165
assert.Error(t, err, "ListPartitions should fail")
158166
}

‎internal/infra/postgresql/table.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
11
package postgresql
22

3-
import "fmt"
3+
import (
4+
"fmt"
5+
6+
"github.com/jackc/pgx/v5"
7+
)
48

59
func (p Postgres) CreateTableLikeTable(schema, table, parent string) error {
6-
query := fmt.Sprintf("CREATE TABLE %s.%s (LIKE %s.%s INCLUDING ALL)", schema, table, schema, parent)
10+
query := fmt.Sprintf("CREATE TABLE %s (LIKE %s INCLUDING ALL)",
11+
pgx.Identifier{schema, table}.Sanitize(),
12+
pgx.Identifier{schema, parent}.Sanitize())
713
p.logger.Debug("Create table", "schema", schema, "table", table, "query", query)
814

915
_, err := p.conn.Exec(p.ctx, query)

‎internal/infra/postgresql/table_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"testing"
77

8+
"github.com/jackc/pgx/v5"
89
"github.com/pashagolub/pgxmock/v3"
910
"github.com/stretchr/testify/assert"
1011
)
@@ -14,7 +15,9 @@ func TestCreateTableLikeTable(t *testing.T) {
1415
table := "my_table"
1516
parentTable := "parent_table"
1617

17-
query := fmt.Sprintf(`CREATE TABLE %s.%s (LIKE %s.%s INCLUDING ALL)`, schema, table, schema, parentTable)
18+
query := fmt.Sprintf("CREATE TABLE %s (LIKE %s INCLUDING ALL)",
19+
pgx.Identifier{schema, table}.Sanitize(),
20+
pgx.Identifier{schema, parentTable}.Sanitize())
1821

1922
mock, p := setupMock(t, pgxmock.QueryMatcherEqual)
2023

‎scripts/bats/30_provisioning.bats‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,5 +474,38 @@ EOF
474474
run execute_sql_commands "$check_query"
475475
assert_output "{i,i,i,i,f}"
476476

477+
rm "$CONFIGURATION_FILE"
478+
}
479+
480+
@test "Test provisioning with special chars in the table name" {
481+
local CONFIGURATION=$(cat << EOF
482+
partitions:
483+
unittest1:
484+
schema: public
485+
table: table's name
486+
interval: daily
487+
partitionKey: created_at
488+
cleanupPolicy: detach
489+
retention: 1
490+
preProvisioned: 1
491+
EOF
492+
)
493+
local CONFIGURATION_FILE=$(generate_configuration_file "${CONFIGURATION}")
494+
495+
create_partitioned_table "\"table's name\""
496+
497+
PPM_WORK_DATE="2025-02-01" run "$PPM_PROG" run provisioning -c ${CONFIGURATION_FILE}
498+
assert_success
499+
500+
local expected=$(cat <<'EOF'
501+
public|table's name_2025_01_31|2025-01-31|2025-02-01
502+
public|table's name_2025_02_01|2025-02-01|2025-02-02
503+
public|table's name_2025_02_02|2025-02-02|2025-02-03
504+
EOF
505+
)
506+
run list_existing_partitions "public" "table's name"
507+
assert_output "$expected"
508+
509+
477510
rm "$CONFIGURATION_FILE"
478511
}

‎scripts/bats/test/libs/sql.bash‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,14 @@ EOSQL
3838

3939
list_existing_partitions() {
4040
# mandatory arguments
41-
local PARENT_SCHEMA=$1
42-
local PARENT_TABLE=$2
41+
local PARENT_SCHEMA="$1"
42+
local PARENT_TABLE="$2"
4343

44-
psql --tuples-only --no-align --quiet --dbname="$PPM_DATABASE" -v parent_schema=${PARENT_SCHEMA} -v parent_table=${PARENT_TABLE} <<'EOSQL'
44+
psql --tuples-only --no-align --quiet --dbname="$PPM_DATABASE" -v parent_schema="${PARENT_SCHEMA}" -v parent_table="${PARENT_TABLE}" <<'EOSQL'
4545
WITH parts as (
4646
SELECT
4747
relnamespace::regnamespace as schema,
48-
c.oid::pg_catalog.regclass AS part_name,
48+
c.relname AS part_name,
4949
regexp_match(pg_get_expr(c.relpartbound, c.oid),
5050
'FOR VALUES FROM \(''(.*)''\) TO \(''(.*)''\)') AS bounds
5151
FROM

0 commit comments

Comments
 (0)