Skip to content

Commit 19e1453

Browse files
exAspArkexAspArk
committed
Allow per-column permissions in runtime. Deprecate PG include_schemas
1 parent 6e20856 commit 19e1453

13 files changed

Lines changed: 217 additions & 185 deletions

‎README.md‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -131,17 +131,6 @@ docker run \
131131
ghcr.io/bemihq/bemidb:latest syncer-postgres
132132
```
133133

134-
To include all tables from a specific schema:
135-
136-
```sh
137-
docker run \
138-
-e SOURCE_POSTGRES_DATABASE_URL=postgres://user:password@host.docker.internal:5432/source \
139-
-e SOURCE_POSTGRES_INCLUDE_SCHEMAS=public \ # A comma-separated list of schemas to include
140-
-e DESTINATION_SCHEMA_NAME=postgres \
141-
-e AWS_REGION -e AWS_S3_BUCKET -e AWS_ACCESS_KEY_ID -e AWS_SECRET_ACCESS_KEY -e CATALOG_DATABASE_URL \
142-
ghcr.io/bemihq/bemidb:latest syncer-postgres
143-
```
144-
145134
To exclude specific tables during the sync:
146135

147136
```sh
@@ -212,7 +201,6 @@ export AWS_S3_ENDPOINT=http://localhost:9000
212201
|-------------------------------------|---------------|----------------------------------------------------------------------|
213202
| `DESTINATION_SCHEMA_NAME` | Required | Schema name in BemiDB to sync data to. |
214203
| `SOURCE_POSTGRES_DATABASE_URL` | Required | Postgres database URL to sync data from. |
215-
| `SOURCE_POSTGRES_INCLUDE_SCHEMAS` | | List of schemas to include in sync. Comma-separated `schema.table`. |
216204
| `SOURCE_POSTGRES_INCLUDE_TABLES` | | List of tables to include in sync. Comma-separated `schema.table`. |
217205
| `SOURCE_POSTGRES_EXCLUDE_TABLES` | | List of tables to exclude from sync. Comma-separated `schema.table`. |
218206

‎src/common/common_config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package common
22

33
const (
4-
VERSION = "1.3.0"
4+
VERSION = "1.4.0"
55

66
ENV_LOG_LEVEL = "BEMIDB_LOG_LEVEL"
77
ENV_DISABLE_ANONYMOUS_ANALYTICS = "BEMIDB_DISABLE_ANONYMOUS_ANALYTICS"

‎src/common/iceberg_catalog.go‎

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ type IcebergSchemaTable struct {
1818
Table string
1919
}
2020

21+
func (schemaTable IcebergSchemaTable) ToArg() string {
22+
return schemaTable.Schema + "." + schemaTable.Table
23+
}
24+
2125
func (schemaTable IcebergSchemaTable) String() string {
2226
return fmt.Sprintf(`"%s"."%s"`, schemaTable.Schema, schemaTable.Table)
2327
}
@@ -47,29 +51,6 @@ func NewIcebergCatalog(config *CommonConfig) *IcebergCatalog {
4751

4852
// Read ----------------------------------------------------------------------------------------------------------------
4953

50-
func (catalog *IcebergCatalog) Schemas() ([]string, error) {
51-
pgClient := catalog.newPostgresClient()
52-
defer pgClient.Close()
53-
54-
ctx := context.Background()
55-
rows, err := pgClient.Query(ctx, "SELECT table_namespace FROM iceberg_tables WHERE table_name NOT LIKE '%"+TEMP_TABLE_SUFFIX_SYNCING+"' AND table_name NOT LIKE '%"+TEMP_TABLE_SUFFIX_DELETING+"' GROUP BY table_namespace")
56-
if err != nil {
57-
return nil, err
58-
}
59-
defer rows.Close()
60-
61-
var schemas []string
62-
for rows.Next() {
63-
var schema string
64-
err := rows.Scan(&schema)
65-
if err != nil {
66-
return nil, err
67-
}
68-
schemas = append(schemas, schema)
69-
}
70-
return schemas, nil
71-
}
72-
7354
func (catalog *IcebergCatalog) SchemaTables() (Set[IcebergSchemaTable], error) {
7455
pgClient := catalog.newPostgresClient()
7556
defer pgClient.Close()
@@ -120,7 +101,6 @@ func (catalog *IcebergCatalog) MaterializedViews() ([]IcebergMaterializedView, e
120101
Table: table,
121102
Definition: definition,
122103
})
123-
124104
}
125105
return materializedViews, nil
126106
}

‎src/server/iceberg_reader.go‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,6 @@ func NewIcebergReader(config *Config, storageS3 *common.StorageS3, icebergCatalo
1818
}
1919
}
2020

21-
func (reader *IcebergReader) Schemas() (icebergSchemas []string, err error) {
22-
return reader.IcebergCatalog.Schemas()
23-
}
24-
2521
func (reader *IcebergReader) SchemaTables() (icebergSchemaTables common.Set[common.IcebergSchemaTable], err error) {
2622
return reader.IcebergCatalog.SchemaTables()
2723
}

‎src/server/parser_table.go‎

Lines changed: 89 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package main
22

33
import (
4+
"strings"
5+
6+
"github.com/BemiHQ/BemiDB/src/common"
47
pgQuery "github.com/pganalyze/pg_query_go/v6"
58
)
69

@@ -42,35 +45,60 @@ func (parser *ParserTable) IsTableFromInformationSchema(qSchemaTable QuerySchema
4245
return qSchemaTable.Schema == PG_SCHEMA_INFORMATION_SCHEMA
4346
}
4447

45-
// public.table -> FROM iceberg_scan('path') table
46-
// schema.table -> FROM iceberg_scan('path') schema_table
47-
func (parser *ParserTable) MakeIcebergTableNode(queryToIcebergTable QueryToIcebergTable) *pgQuery.Node {
48-
node := pgQuery.MakeSimpleRangeFunctionNode([]*pgQuery.Node{
49-
pgQuery.MakeListNode([]*pgQuery.Node{
50-
pgQuery.MakeFuncCallNode(
51-
[]*pgQuery.Node{
52-
pgQuery.MakeStrNode("iceberg_scan"),
53-
},
54-
[]*pgQuery.Node{
55-
pgQuery.MakeAConstStrNode(
56-
queryToIcebergTable.IcebergTablePath,
57-
0,
58-
),
59-
},
60-
0,
61-
),
62-
}),
63-
})
64-
65-
// DuckDB doesn't support aliases on iceberg_scan() functions, so we need to wrap it in a nested select that can have an alias
66-
selectStarNode := pgQuery.MakeResTargetNodeWithVal(
67-
pgQuery.MakeColumnRefNode(
68-
[]*pgQuery.Node{pgQuery.MakeAStarNode()},
69-
0,
70-
),
71-
0,
72-
)
73-
return parser.utils.MakeSubselectFromNode(queryToIcebergTable.QuerySchemaTable, []*pgQuery.Node{selectStarNode}, node)
48+
// public.table -> (SELECT * FROM iceberg_scan('path')) table
49+
// schema.table -> (SELECT * FROM iceberg_scan('path')) schema_table
50+
// public.table -> (SELECT permitted, columns FROM iceberg_scan('path')) table
51+
// public.table -> (SELECT NULL WHERE FALSE) table
52+
// public.table t -> (SELECT * FROM iceberg_scan('path')) t
53+
func (parser *ParserTable) MakeIcebergTableNode(queryToIcebergTable QueryToIcebergTable, permissions *map[string][]string) *pgQuery.Node {
54+
var query string
55+
if permissions == nil {
56+
query = "SELECT * FROM iceberg_scan('" + queryToIcebergTable.IcebergTablePath + "')"
57+
} else if columnNames, allowed := (*permissions)[queryToIcebergTable.QuerySchemaTable.ToIcebergSchemaTable().ToArg()]; allowed {
58+
query = "SELECT " + strings.Join(columnNames, ", ") + " FROM iceberg_scan('" + queryToIcebergTable.IcebergTablePath + "')"
59+
} else {
60+
query = "SELECT NULL WHERE FALSE"
61+
}
62+
63+
return parser.makeSubselectNode(query, queryToIcebergTable.QuerySchemaTable)
64+
}
65+
66+
// information_schema.tables -> (SELECT * FROM main.tables) information_schema_tables
67+
// information_schema.tables -> (SELECT * FROM main.tables WHERE table_schema || '.' || table_name IN ('permitted.table')) information_schema_tables
68+
// information_schema.tables t -> (SELECT * FROM main.tables) t
69+
func (parser *ParserTable) MakeInformationSchemaTablesNode(qSchemaTable QuerySchemaTable, permissions *map[string][]string) *pgQuery.Node {
70+
query := "SELECT * FROM main.tables"
71+
72+
if permissions != nil {
73+
quotedSchemaTableNames := []string{}
74+
for schemaTable := range *permissions {
75+
quotedSchemaTableNames = append(quotedSchemaTableNames, "'"+schemaTable+"'")
76+
}
77+
query += " WHERE table_schema || '.' || table_name IN (" + strings.Join(quotedSchemaTableNames, ", ") + ")"
78+
}
79+
80+
return parser.makeSubselectNode(query, qSchemaTable)
81+
}
82+
83+
// information_schema.columns -> (SELECT * FROM main.columns) information_schema_columns
84+
// information_schema.columns -> (SELECT * FROM main.columns WHERE (table_schema || '.' || table_name IN ('permitted.table') AND column_name IN ('permitted', 'columns')) OR ...) information_schema_columns
85+
// information_schema.columns c -> (SELECT * FROM main.columns) c
86+
func (parser *ParserTable) MakeInformationSchemaColumnsNode(qSchemaTable QuerySchemaTable, permissions *map[string][]string) *pgQuery.Node {
87+
query := "SELECT * FROM main.columns"
88+
89+
if permissions != nil {
90+
conditions := []string{}
91+
for schemaTable, columnNames := range *permissions {
92+
quotedColumnNames := []string{}
93+
for _, columnName := range columnNames {
94+
quotedColumnNames = append(quotedColumnNames, "'"+columnName+"'")
95+
}
96+
conditions = append(conditions, "(table_schema || '.' || table_name = '"+schemaTable+"' AND column_name IN ("+strings.Join(quotedColumnNames, ", ")+"))")
97+
}
98+
query += " WHERE " + strings.Join(conditions, " OR ")
99+
}
100+
101+
return parser.makeSubselectNode(query, qSchemaTable)
74102
}
75103

76104
func (parser *ParserTable) TopLevelSchemaFunction(rangeFunction *pgQuery.RangeFunction) *QuerySchemaFunction {
@@ -108,3 +136,35 @@ func (parser *ParserTable) SetAliasIfNotExists(rangeFunction *pgQuery.RangeFunct
108136

109137
rangeFunction.Alias = &pgQuery.Alias{Aliasname: alias}
110138
}
139+
140+
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
141+
142+
// (query) AS qSchemaTable
143+
func (parser *ParserTable) makeSubselectNode(query string, qSchemaTable QuerySchemaTable) *pgQuery.Node {
144+
queryTree, err := pgQuery.Parse(query)
145+
common.PanicIfError(parser.config.CommonConfig, err)
146+
147+
alias := qSchemaTable.Alias
148+
if alias == "" {
149+
if qSchemaTable.Schema == PG_SCHEMA_PUBLIC || qSchemaTable.Schema == "" {
150+
alias = qSchemaTable.Table
151+
} else {
152+
alias = qSchemaTable.Schema + "_" + qSchemaTable.Table
153+
}
154+
}
155+
156+
return &pgQuery.Node{
157+
Node: &pgQuery.Node_RangeSubselect{
158+
RangeSubselect: &pgQuery.RangeSubselect{
159+
Subquery: &pgQuery.Node{
160+
Node: &pgQuery.Node_SelectStmt{
161+
SelectStmt: queryTree.Stmts[0].Stmt.GetSelectStmt(),
162+
},
163+
},
164+
Alias: &pgQuery.Alias{
165+
Aliasname: alias,
166+
},
167+
},
168+
},
169+
}
170+
}

‎src/server/parser_utils.go‎

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -32,35 +32,6 @@ func (utils *ParserUtils) SchemaFunction(functionCall *pgQuery.FuncCall) *QueryS
3232
}
3333
}
3434

35-
func (utils *ParserUtils) MakeSubselectFromNode(qSchemaTable QuerySchemaTable, targetList []*pgQuery.Node, fromNode *pgQuery.Node) *pgQuery.Node {
36-
alias := qSchemaTable.Alias
37-
if alias == "" {
38-
if qSchemaTable.Schema == PG_SCHEMA_PUBLIC || qSchemaTable.Schema == "" {
39-
alias = qSchemaTable.Table
40-
} else {
41-
alias = qSchemaTable.Schema + "_" + qSchemaTable.Table
42-
}
43-
}
44-
45-
return &pgQuery.Node{
46-
Node: &pgQuery.Node_RangeSubselect{
47-
RangeSubselect: &pgQuery.RangeSubselect{
48-
Subquery: &pgQuery.Node{
49-
Node: &pgQuery.Node_SelectStmt{
50-
SelectStmt: &pgQuery.SelectStmt{
51-
TargetList: targetList,
52-
FromClause: []*pgQuery.Node{fromNode},
53-
},
54-
},
55-
},
56-
Alias: &pgQuery.Alias{
57-
Aliasname: alias,
58-
},
59-
},
60-
},
61-
}
62-
}
63-
6435
func (utils *ParserUtils) MakeNullNode() *pgQuery.Node {
6536
return &pgQuery.Node{
6637
Node: &pgQuery.Node_AConst{

‎src/server/pg_constants.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ const (
1717
PG_TABLE_PG_CLASS = "pg_class"
1818
PG_TABLE_PG_STAT_USER_TABLES = "pg_stat_user_tables"
1919
PG_TABLE_TABLES = "tables"
20+
PG_TABLE_COLUMNS = "columns"
2021

2122
PG_VAR_SEARCH_PATH = "search_path"
2223
)

‎src/server/query_handler.go‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,7 @@ import (
1515
)
1616

1717
const (
18-
FALLBACK_SQL_QUERY = "SELECT 1"
19-
INSPECT_SQL_COMMENT = " --INSPECT"
18+
FALLBACK_SQL_QUERY = "SELECT 1"
2019
)
2120

2221
type QueryHandler struct {
@@ -59,8 +58,6 @@ func NewQueryHandler(config *Config, serverDuckdbClient *common.DuckdbClient) *Q
5958
ResponseHandler: NewResponseHandler(config),
6059
}
6160

62-
queryHandler.createSchemas(icebergReader)
63-
6461
return queryHandler
6562
}
6663

@@ -238,21 +235,6 @@ func (queryHandler *QueryHandler) HandleExecuteQuery(message *pgproto3.Execute,
238235
return queryHandler.rowsToDataMessages(preparedStatement.Rows, preparedStatement.OriginalQuery)
239236
}
240237

241-
func (queryHandler *QueryHandler) createSchemas(icebergReader *IcebergReader) {
242-
ctx := context.Background()
243-
schemas, err := icebergReader.Schemas()
244-
common.PanicIfError(queryHandler.Config.CommonConfig, err)
245-
246-
for _, schema := range schemas {
247-
_, err := queryHandler.ServerDuckdbClient.ExecContext(
248-
ctx,
249-
"CREATE SCHEMA IF NOT EXISTS \"$schema\"",
250-
map[string]string{"schema": schema},
251-
)
252-
common.PanicIfError(queryHandler.Config.CommonConfig, err)
253-
}
254-
}
255-
256238
func (queryHandler *QueryHandler) rowsToDescriptionMessages(rows *sql.Rows, originalQuery string) ([]pgproto3.Message, error) {
257239
cols, err := rows.ColumnTypes()
258240
if err != nil {

‎src/server/query_handler_test.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -375,6 +375,11 @@ func TestHandleQuery(t *testing.T) {
375375
"types": {uint32ToString(pgtype.TextOID)},
376376
"values": {"memory", "postgres", "test_table", "BASE TABLE", "", "", "", "", "", "YES", "NO", ""},
377377
},
378+
"SELECT * FROM information_schema.tables WHERE table_schema = 'postgres' /*BEMIDB_PERMISSIONS {\"postgres.test_table\": [\"id\"]} BEMIDB_PERMISSIONS*/": {
379+
"description": {"table_catalog", "table_schema", "table_name", "table_type", "self_referencing_column_name", "reference_generation", "user_defined_type_catalog", "user_defined_type_schema", "user_defined_type_name", "is_insertable_into", "is_typed", "commit_action"},
380+
"types": {uint32ToString(pgtype.TextOID)},
381+
"values": {"memory", "postgres", "test_table", "BASE TABLE", "", "", "", "", "", "YES", "NO", ""},
382+
},
378383
// information_schema.columns
379384
"SELECT udt_name FROM information_schema.columns WHERE table_schema = 'postgres' AND table_name = 'test_table' AND column_name = 'id'": {
380385
"description": {"udt_name"},
@@ -576,6 +581,11 @@ func TestHandleQuery(t *testing.T) {
576581
"types": {uint32ToString(pgtype.TextOID)},
577582
"values": {"text"},
578583
},
584+
"SELECT udt_name FROM information_schema.columns WHERE table_schema = 'postgres' AND table_name = 'test_table' /*BEMIDB_PERMISSIONS {\"postgres.test_table\": [\"id\"]} BEMIDB_PERMISSIONS*/": {
585+
"description": {"udt_name"},
586+
"types": {uint32ToString(pgtype.TextOID)},
587+
"values": {"int4"},
588+
},
579589
})
580590
})
581591

@@ -1425,6 +1435,25 @@ func TestHandleQuery(t *testing.T) {
14251435
}
14261436
})
14271437

1438+
t.Run("Returns an error if permission for a column is denied", func(t *testing.T) {
1439+
_, err := queryHandler.HandleSimpleQuery("SELECT id, bit_column FROM postgres.test_table /*BEMIDB_PERMISSIONS {\"postgres.test_table\": [\"id\"]} BEMIDB_PERMISSIONS*/")
1440+
1441+
if err == nil {
1442+
t.Errorf("Expected an error, got nil")
1443+
}
1444+
1445+
expectedErrorMessage := strings.Join([]string{
1446+
"Binder Error: Referenced column \"bit_column\" not found in FROM clause!",
1447+
"Candidate bindings: \"id\"",
1448+
"",
1449+
"LINE 1: SELECT id, bit_column FROM (SELECT id FROM iceberg_scan('s3://bemidb...",
1450+
" ^",
1451+
}, "\n")
1452+
if err.Error() != expectedErrorMessage {
1453+
t.Errorf("Expected the error to be '"+expectedErrorMessage+"', got %v", err.Error())
1454+
}
1455+
})
1456+
14281457
t.Run("Returns a result without a row description for SET queries", func(t *testing.T) {
14291458
messages, err := queryHandler.HandleSimpleQuery("SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ UNCOMMITTED")
14301459

0 commit comments

Comments
 (0)