Skip to content

Commit 58652bf

Browse files
exAspArkexAspArk
committed
Add support for partitioned tables
1 parent 7c956ba commit 58652bf

16 files changed

Lines changed: 389 additions & 302 deletions

‎src/config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import (
88
)
99

1010
const (
11-
VERSION = "0.49.3"
11+
VERSION = "0.50.0"
1212

1313
ENV_PORT = "BEMIDB_PORT"
1414
ENV_DATABASE = "BEMIDB_DATABASE"

‎src/custom_types.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,14 @@ func (pgSchemaTable PgSchemaTable) String() string {
132132
return fmt.Sprintf(`"%s"."%s"`, pgSchemaTable.Schema, pgSchemaTable.Table)
133133
}
134134

135+
func (pgSchemaTable PgSchemaTable) ParentPartitionedTableString() string {
136+
if pgSchemaTable.ParentPartitionedTable == "" {
137+
return ""
138+
}
139+
140+
return fmt.Sprintf(`"%s"."%s"`, pgSchemaTable.Schema, pgSchemaTable.ParentPartitionedTable)
141+
}
142+
135143
func (pgSchemaTable PgSchemaTable) ToConfigArg() string {
136144
return fmt.Sprintf(`%s.%s`, pgSchemaTable.Schema, pgSchemaTable.Table)
137145
}

‎src/iceberg_reader.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ func (reader *IcebergReader) MetadataFilePath(icebergSchemaTable IcebergSchemaTa
3333
return reader.storage.IcebergMetadataFilePath(icebergSchemaTable)
3434
}
3535

36+
func (reader *IcebergReader) InternalTableMetadata(pgSchemaTable PgSchemaTable) (internalTableMetadata InternalTableMetadata, err error) {
37+
return reader.storage.InternalTableMetadata(pgSchemaTable)
38+
}
39+
3640
func (reader *IcebergReader) InternalStartSqlFile() io.ReadCloser {
3741
return reader.storage.InternalStartSqlFile()
3842
}

‎src/iceberg_writer.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package main
2+
3+
type IcebergWriter struct {
4+
config *Config
5+
storage StorageInterface
6+
}
7+
8+
func NewIcebergWriter(config *Config) *IcebergWriter {
9+
storage := NewStorage(config)
10+
return &IcebergWriter{config: config, storage: storage}
11+
}
12+
13+
func (writer *IcebergWriter) DeleteSchema(icebergSchema string) (err error) {
14+
return writer.storage.DeleteSchema(icebergSchema)
15+
}
16+
17+
func (writer *IcebergWriter) DeleteSchemaTable(icebergSchemaTable IcebergSchemaTable) (err error) {
18+
return writer.storage.DeleteSchemaTable(icebergSchemaTable)
19+
}
20+
21+
func (writer *IcebergWriter) WriteInternalStartSqlFile(queries []string) (err error) {
22+
return writer.storage.WriteInternalStartSqlFile(queries)
23+
}

‎src/init_test.go‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"os"
66
)
77

8-
var PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
8+
var PUBLIC_SCHEMA_TEST_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
99
{
1010
ColumnName: "id",
1111
DataType: "integer",
@@ -283,8 +283,7 @@ var PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
283283
Namespace: "public",
284284
},
285285
}
286-
287-
var PUBLIC_TEST_TABLE_LOADED_ROWS = [][]string{
286+
var PUBLIC_SCHEMA_TEST_TABLE_LOADED_ROWS = [][]string{
288287
{
289288
"1", // id
290289
"1", // bit_column
@@ -373,7 +372,7 @@ var PUBLIC_TEST_TABLE_LOADED_ROWS = [][]string{
373372
},
374373
}
375374

376-
var TEST_SCHEMA_SIMPLE_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
375+
var TEST_SCHEMA_EMPTY_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
377376
{
378377
ColumnName: "id",
379378
DataType: "integer",
@@ -387,16 +386,33 @@ var TEST_SCHEMA_SIMPLE_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
387386
Namespace: "pg_catalog",
388387
},
389388
}
389+
var TEST_SCHEMA_EMPTY_TABLE_LOADED_ROWS = [][]string{}
390390

391-
var TEST_SCHEMA_SIMPLE_TABLE_LOADED_ROWS = [][]string{}
391+
var PUBLIC_SCHEMA_PARTITIONED_TABLE_PG_SCHEMA_COLUMNS = []PgSchemaColumn{
392+
{
393+
ColumnName: "timestamp_column",
394+
DataType: "timestamp without time zone",
395+
UdtName: "timestamp",
396+
IsNullable: "NO",
397+
OrdinalPosition: "1",
398+
CharacterMaximumLength: "0",
399+
NumericPrecision: "0",
400+
NumericScale: "0",
401+
DatetimePrecision: "6",
402+
Namespace: "pg_catalog",
403+
},
404+
}
405+
var PUBLIC_SCHEMA_PARTITIONED_TABLE1_LOADED_ROWS = [][]string{{"2024-01-01 01:02:03.123456"}}
406+
var PUBLIC_SCHEMA_PARTITIONED_TABLE2_LOADED_ROWS = [][]string{{"2024-02-12 11:12:13"}}
407+
var PUBLIC_SCHEMA_PARTITIONED_TABLE3_LOADED_ROWS = [][]string{{"2024-03-30 23:59:59"}}
392408

393409
func init() {
394410
loadTestConfig()
395411

396-
for i := range PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS {
397-
PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS[i].OrdinalPosition = IntToString(i + 1)
398-
if PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS[i].IsNullable == "" {
399-
PUBLIC_TEST_TABLE_PG_SCHEMA_COLUMNS[i].IsNullable = "YES"
412+
for i := range PUBLIC_SCHEMA_TEST_TABLE_PG_SCHEMA_COLUMNS {
413+
PUBLIC_SCHEMA_TEST_TABLE_PG_SCHEMA_COLUMNS[i].OrdinalPosition = IntToString(i + 1)
414+
if PUBLIC_SCHEMA_TEST_TABLE_PG_SCHEMA_COLUMNS[i].IsNullable == "" {
415+
PUBLIC_SCHEMA_TEST_TABLE_PG_SCHEMA_COLUMNS[i].IsNullable = "YES"
400416
}
401417
}
402418
}
@@ -411,6 +427,7 @@ func loadTestConfig() *Config {
411427
config.StoragePath = "../iceberg-test"
412428
config.LogLevel = "ERROR"
413429
config.DisableAnonymousAnalytics = true
430+
config.Pg.DatabaseUrl = "postgresql://postgres:postgres@localhost:5432/dbname"
414431

415432
return config
416433
}

‎src/parser_table.go‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@ import (
44
pgQuery "github.com/pganalyze/pg_query_go/v5"
55
)
66

7+
type QueryToIcebergTable struct {
8+
QuerySchemaTable QuerySchemaTable
9+
IcebergTablePath string
10+
}
11+
712
type ParserTable struct {
813
config *Config
914
utils *ParserUtils
@@ -39,7 +44,7 @@ func (parser *ParserTable) IsTableFromInformationSchema(qSchemaTable QuerySchema
3944

4045
// public.table -> FROM iceberg_scan('path', skip_schema_inference = true) table
4146
// schema.table -> FROM iceberg_scan('path', skip_schema_inference = true) schema_table
42-
func (parser *ParserTable) MakeIcebergTableNode(tablePath string, qSchemaTable QuerySchemaTable) *pgQuery.Node {
47+
func (parser *ParserTable) MakeIcebergTableNode(queryToIcebergTable QueryToIcebergTable) *pgQuery.Node {
4348
node := pgQuery.MakeSimpleRangeFunctionNode([]*pgQuery.Node{
4449
pgQuery.MakeListNode([]*pgQuery.Node{
4550
pgQuery.MakeFuncCallNode(
@@ -48,7 +53,7 @@ func (parser *ParserTable) MakeIcebergTableNode(tablePath string, qSchemaTable Q
4853
},
4954
[]*pgQuery.Node{
5055
pgQuery.MakeAConstStrNode(
51-
tablePath,
56+
queryToIcebergTable.IcebergTablePath,
5257
0,
5358
),
5459
pgQuery.MakeAExprNode(
@@ -72,7 +77,7 @@ func (parser *ParserTable) MakeIcebergTableNode(tablePath string, qSchemaTable Q
7277
),
7378
0,
7479
)
75-
return parser.utils.MakeSubselectFromNode(qSchemaTable, []*pgQuery.Node{selectStarNode}, node)
80+
return parser.utils.MakeSubselectFromNode(queryToIcebergTable.QuerySchemaTable, []*pgQuery.Node{selectStarNode}, node)
7681
}
7782

7883
func (parser *ParserTable) TopLevelSchemaFunction(rangeFunction *pgQuery.RangeFunction) *QuerySchemaFunction {

‎src/query_handler.go‎

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import (
1414
"github.com/jackc/pgx/v5/pgproto3"
1515
"github.com/jackc/pgx/v5/pgtype"
1616
duckDb "github.com/marcboeker/go-duckdb"
17-
pgQuery "github.com/pganalyze/pg_query_go/v5"
1817
)
1918

2019
const (
@@ -259,7 +258,7 @@ func NewQueryHandler(config *Config, duckdb *Duckdb, icebergReader *IcebergReade
259258
}
260259

261260
func (queryHandler *QueryHandler) HandleSimpleQuery(originalQuery string) ([]pgproto3.Message, error) {
262-
queryStatements, originalQueryStatements, err := queryHandler.parseAndRemapQuery(originalQuery)
261+
queryStatements, originalQueryStatements, err := queryHandler.queryRemapper.ParseAndRemapQuery(originalQuery)
263262
if err != nil {
264263
return nil, err
265264
}
@@ -309,7 +308,7 @@ func (queryHandler *QueryHandler) HandleSimpleQuery(originalQuery string) ([]pgp
309308
func (queryHandler *QueryHandler) HandleParseQuery(message *pgproto3.Parse) ([]pgproto3.Message, *PreparedStatement, error) {
310309
ctx := context.Background()
311310
originalQuery := string(message.Query)
312-
queryStatements, _, err := queryHandler.parseAndRemapQuery(originalQuery)
311+
queryStatements, _, err := queryHandler.queryRemapper.ParseAndRemapQuery(originalQuery)
313312
if err != nil {
314313
return nil, nil, err
315314
}
@@ -495,42 +494,6 @@ func (queryHandler *QueryHandler) rowsToDataMessages(rows *sql.Rows, originalQue
495494
return messages, nil
496495
}
497496

498-
func (queryHandler *QueryHandler) parseAndRemapQuery(query string) ([]string, []string, error) {
499-
queryTree, err := pgQuery.Parse(query)
500-
if err != nil {
501-
return nil, nil, fmt.Errorf("couldn't parse query: %s. %w", query, err)
502-
}
503-
504-
if strings.HasSuffix(query, INSPECT_SQL_COMMENT) {
505-
LogDebug(queryHandler.config, queryTree.Stmts)
506-
}
507-
508-
var originalQueryStatements []string
509-
for _, stmt := range queryTree.Stmts {
510-
originalQueryStatement, err := pgQuery.Deparse(&pgQuery.ParseResult{Stmts: []*pgQuery.RawStmt{stmt}})
511-
if err != nil {
512-
return nil, nil, fmt.Errorf("couldn't deparse query: %s. %w", query, err)
513-
}
514-
originalQueryStatements = append(originalQueryStatements, originalQueryStatement)
515-
}
516-
517-
remappedStatements, err := queryHandler.queryRemapper.RemapStatements(queryTree.Stmts)
518-
if err != nil {
519-
return nil, nil, fmt.Errorf("couldn't remap query: %s. %w", query, err)
520-
}
521-
522-
var queryStatements []string
523-
for _, remappedStatement := range remappedStatements {
524-
queryStatement, err := pgQuery.Deparse(&pgQuery.ParseResult{Stmts: []*pgQuery.RawStmt{remappedStatement}})
525-
if err != nil {
526-
return nil, nil, fmt.Errorf("couldn't deparse remapped query: %s. %w", query, err)
527-
}
528-
queryStatements = append(queryStatements, queryStatement)
529-
}
530-
531-
return queryStatements, originalQueryStatements, nil
532-
}
533-
534497
func (queryHandler *QueryHandler) generateRowDescription(cols []*sql.ColumnType) *pgproto3.RowDescription {
535498
description := pgproto3.RowDescription{Fields: []pgproto3.FieldDescription{}}
536499

0 commit comments

Comments
 (0)