feat(rest): support query auth for REST catalog tables (#534)
diff --git a/docs/source/user_guide/catalog.rst b/docs/source/user_guide/catalog.rst index 0592a7b..f4a46f7 100644 --- a/docs/source/user_guide/catalog.rst +++ b/docs/source/user_guide/catalog.rst
@@ -319,10 +319,27 @@ Reading through the catalog ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -To read a native table with the per-table temporary credentials a catalog issues, -pass the catalog and table identifier to ``ReadContextBuilder::WithCatalog`` or -``ScanContextBuilder::WithCatalog``. This is a shorthand for asking the catalog -for that table's file system and passing it in through ``WithFileSystem``: +To read a native table through its catalog, pass the catalog and table +identifier to ``ReadContextBuilder::WithCatalog`` or +``ScanContextBuilder::WithCatalog``. The catalog then supplies: + +* the table's file system, with the per-table temporary credentials the catalog + issues, unless ``WithFileSystem`` sets one; +* whether the table enables ``query-auth.enabled``, read from the schema the + catalog holds for the branch read; +* for a table that does, its row filters and column masks. + +Passing only the file system through ``WithFileSystem`` gives neither of the +last two. + +The fields a scan plans with and a read reads with do not come from the catalog, +but from the schema ``SetTableSchema`` hands in for the main branch or, otherwise, +from the latest schema published under the table path of the branch read. That +schema has to exist even when the catalog holds the current one: a table whose +schema only the catalog holds cannot be scanned or read this way yet, and when +the two differ, the fields are those of the schema under the table path. A +query-auth scan also checks column identities and rule types against the +catalog's schema on every plan: .. code-block:: cpp @@ -347,19 +364,253 @@ system through ``Catalog::GetTableFileSystem`` and uses it as-is for the schema, snapshots, manifests and data under the table path, so it signs every access with the table's credentials and reloads them as they expire. An explicit -``WithFileSystem`` takes precedence, so the catalog is not asked. A catalog that -issues no per-table credentials returns its catalog-level file system, so this is -also how you read with the catalog's own object-store credentials. For a format -table, use ``ReadContextBuilder(FormatTable)`` or ``ScanContextBuilder(FormatTable)`` -instead; ``WithCatalog`` is rejected because the table already carries the file -system it was loaded through. +``WithFileSystem`` takes precedence, so the catalog is not asked for one. A +catalog that issues no per-table credentials returns its catalog-level file +system, so this is also how you read with the catalog's own object-store +credentials. For a format table, use ``ReadContextBuilder(FormatTable)`` or +``ScanContextBuilder(FormatTable)`` instead; ``WithCatalog`` is rejected because +the table already carries the file system it was loaded through. + +Row filters and column masks +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +A table with ``query-auth.enabled = true`` lets the catalog restrict what a read +returns with row filters and column masks, as in Java. The option is enabled when +any of these sets it: + +* the latest schema under the table path; +* the options of the scan or read; +* for a scan or read built ``WithCatalog``, the schema the catalog holds for the + branch read, which ``TableScan::Create`` and ``TableRead::Create`` load as Java + loads the table of that branch. + +A schema handed in with ``SetTableSchema`` for the main branch does not decide it +alone: unless the catalog already enabled the option, the latest schema under the +table path is checked too, so a cache taken before the option was turned on still +applies the rules. + +What the rules may contain: + +* Row filters support the ``AND`` and ``OR`` of the comparison, ``IN``, + ``BETWEEN``, null, prefix, suffix, ``CONTAINS`` and ``LIKE`` predicates on a + column referenced with ``FIELD_REF``, and the ``TRUE`` and ``FALSE`` constants + Java's ``PredicateBuilder.alwaysTrue()`` and ``alwaysFalse()`` write on a + ``NULL`` transform; a ``FALSE`` filter reads no row. A predicate on any other + transform, such as ``UPPER(name)``, is refused. ``LIKE`` matches as Java's + does: ``_`` is any one character, a line terminator included. +* Their literals may be ``BOOLEAN``, ``TINYINT``, ``SMALLINT``, ``INT``, + ``BIGINT``, ``FLOAT``, ``DOUBLE``, ``STRING``, ``DECIMAL``, ``DATE``, + ``TIMESTAMP`` and ``TIMESTAMP_LTZ``, the temporal ones in the ISO text format + written by Java's ``LocalDate``, ``LocalDateTime`` and ``Instant``, with an + instant ending in ``Z``. A year outside + ``0000`` to ``9999`` takes a sign as in Java, a fraction at most nine digits, + and a timestamp must lie within the range of epoch milliseconds. A filter on a + column of another type, such as ``TIME``, is refused with ``NotImplemented``; + a ``FIELD_REF`` mask still copies such a column. ``DOUBLE`` literals use + full-precision parsing to preserve Java's comparison boundaries; non-finite + floating-point literals are refused. Integer ``FLOAT`` literals preserve Java's + direct integer-to-float rounding, including integers beyond 64 bits. +* Column masks support the ``NULL``, ``FIELD_REF``, ``CONCAT``, ``CONCAT_WS`` and + ``SUBSTRING`` transforms, whose inputs are literals, field references or null; + a transform nested in the input of another is refused. As in Java, + ``SUBSTRING`` returns null if any input is null, including its position or + length. +* A mask of a system column such as ``_ROW_ID``, which a query-auth read never + projects, is parsed but not applied, as Java ignores one the query does not + project. If the table schema contains an actual column with such a name, the + query fails instead of silently ignoring that column's mask. +* Malformed rules fail the query, including invalid UTF-8 and unescaped NUL + bytes. The ``literals`` array of an ``IN`` or ``NOT_IN`` rule must be present + and non-null; an explicit ``[]`` remains valid. A rule, or the auth response + carrying it, that names a JSON member twice is refused rather than read as + either value. An escaped NUL in a string literal or column name is preserved. + +What would not apply the rules fails with ``NotImplemented``: search, an explicit +read schema (``SetReadSchema``, use ``SetReadFieldNames`` instead), a projected +or filtered column outside the table schema such as ``_ROW_ID``, real-time +reads, ``GlobalIndexScan::Create``, a partition filter on a masked partition +key, a count reader, the ``$files``, ``$manifests``, ``$binlog``, ``$audit_log`` +and ``$ro`` system tables, fallback reads of a query-auth table, and a fallback +branch that enables the option on a table that does not, including when only its +catalog schema or a newer schema than the cached one enables it. Each fallback +branch needs its own authorized splits and schema validation. Passing +``query-auth.enabled = false`` as a read option fails with ``Invalid`` when the +table enables the option, and changes nothing otherwise. Writing to and +committing to the table are not affected. A format table refuses the option, +whether its schema or the options of a call set it. ``sys.partitions`` leaves +out the partitions of a query-auth table, whose values its rules may hide, +including when only its latest published schema enables the option and the +catalog's schema has not caught up. +For REST tables, the catalog schema and location come from the same response, +so a concurrent table recreation cannot pair an old policy with a new location. + +Planning and reading a query-auth table +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Every plan of such a table, the first included, asks the catalog for the rules +of the query (``POST /v1/{prefix}/databases/{database}/tables/{table}/auth``), +checks that the columns they name still exist with the types they were written +for and are columns of the schema the scan plans with, and attaches them to each +split. A scan therefore has to be built ``WithCatalog``, and a catalog that does +not support the query auth interface is refused with ``NotImplemented``. The +scan path must be the location ``Catalog::GetTableLocation`` returns for the +table and branch the rules are asked for, so the rules of one table are never +applied to the files of another; a different path fails with ``Invalid`` when +the scan is created or, if the catalog has since bound the table elsewhere, when +it plans. A scan or read whose catalog leaves ``query-auth.enabled`` off is held +to the same location, so that setting is never taken for another table, and a +system table path such as ``tbl$branch_dev$files`` takes it from the branch it +names. An identifier in the ``sys`` database can only be used with a global +system table path (``<warehouse>/sys/<table>``); it cannot skip the catalog policy +of a data table or its per-table system tables. +The branch the identifier, the ``branch`` option, ``WithBranch()`` or a +system table path names must agree, which ``Finish()`` checks before it asks the +catalog for anything. A branch named only by the option or the path is carried +into the identifier before requesting the table's file system, so the +``query-auth.enabled`` setting, the credentials and the rules all come from the +branch the scan plans. +``Finish()`` also refuses, as for a write, a branch the catalog would read as +another object, such as ``MAIN`` or a name holding a ``$`` +(see `Naming the branch`_). The request selects every column unless +``ScanContextBuilder::SetReadFieldNames`` names the columns the query reads; the +columns of its predicate and partition filter are then selected as well. The +catalog authorizes these columns by their current names, so every plan also +checks that each is still the field, by id, that the scan was created with, in +both the latest published schema and the catalog's current schema. The rules +must fit both schemas as well. A mismatch fails the plan, even when the catalog +returns no rules, so a catalog-only rename or column replacement cannot expose +old data through a stale local schema. Refresh the published or cached schema +and create the scan again after such a change. For an external table, which Java +reads without rules, the REST catalog does not ask the server for any once the +table it loads reports itself as external; a failed load or a failed request +fails the plan. + +Each plan asks for the rules again, but the scan pushes the row filter, and the +predicate conjuncts on columns not masked, into the planning of its first plan. +A later plan whose row filter differs, or that masks a column of the predicate +the first one did not, fails with ``Invalid`` rather than plan with filters +pushed down under the old rules. Java can replace its row-filter pruning at each +plan, but likewise refuses a new mask on a column whose user predicate was +already pushed down. Create the paimon-cpp scan again to apply the new rules; to +resume a streaming scan, start the new one from the snapshot after the last one +read, with ``scan.mode = from-snapshot`` and ``scan.snapshot-id``. + +Each split records the table location, branch and columns its rules were +authorized for. Columns are bound by name and field id in the schema the scan +planned with, including every column when the scan selected them all. The split +also records the field id of every column its rules reference. The read applies +the rules a split carries and refuses a split planned without them. It refuses +a split authorized for another table location or branch, even when their schemas +and data file names match. This check also applies to a worker without the +catalog. The read only has to know that the table enables query auth, which the +schema under the table path, ``WithCatalog`` or the read option +``query-auth.enabled = true`` tells it. It refuses a split when a column it +outputs or filters on is not among those the split was authorized for or is now +another field under the same name, as after a rename or a column added since the +plan, and when a column the rules reference was renamed, replaced or dropped +since. It compares names and field ids only: the types the rules declare are +checked against the schema read, and a rule declaring none, such as a ``NULL`` +mask, only needs its column to exist. Plan again after such a schema change. +Its rows go through the row filter, then the column masks, then the read's own +predicate, so that predicate sees the masked values: all of it when +``EnablePredicateFilter`` asks for it, and otherwise, as in Java, its conjuncts +on masked columns, which no scan prunes by. The columns the rules and the +predicate read are read as well and dropped from the output. A masked column is +output as nullable and decoded. A column a mask, the row filter or that +predicate reads is decoded for it if the file reader kept it dictionary-encoded +or as a large string, and is output decoded too unless only a mask reads it. +Under a row filter no column or nested child is output dictionary-encoded, since +a dictionary would still hold the values of the rows the filter drops; otherwise +any other column is output as the file reader returned it. +Neither the predicate of a read nor the conjuncts of a scan's predicate on masked +columns prune files, whose statistics hold raw values, and a scan whose predicate +has such conjuncts does not push its limit down; nor does one whose row filter +reads data columns. A mask reading another masked column and a rule paimon-cpp +does not know fail the query. + +Query-auth splits across runtimes +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Plan and read a table with ``query-auth.enabled = true`` in the same runtime. +paimon-cpp can deserialize supported Java query-auth splits but refuses to read +them. Java rejects the envelope of a split planned by paimon-cpp when +deserializing it. + +A split planned by paimon-cpp serializes in an envelope of its own, the Java +``SplitSerializer`` bytes of the ``QueryAuthSplit`` followed by that scope, which +Java rejects as a corrupted split rather than reading it without the scope. A +``QueryAuthSplit`` Java serialized does not record which columns its rules were +authorized for, so paimon-cpp deserializes it with the same rules but refuses to +read it: its rules may be those of other columns than the read outputs. Written +back, it takes the Java format, though not necessarily the bytes Java wrote: a +null rule list or map becomes an empty one and the masks are ordered by column. +Only a Java split of ``SplitSerializer`` version 1 wrapping a ``DataSplit`` or an +``IndexedSplit``, whose rules paimon-cpp supports, deserializes; one whose +``authResult`` is null is refused. + +Query-auth examples +^^^^^^^^^^^^^^^^^^^ +These examples use ``shared_catalog`` from `Reading through the catalog`_. + +.. code-block:: cpp + + using namespace paimon; + + // Plan with the columns and the filter the read uses. + ScanContextBuilder scan_builder(table_path); + scan_builder.WithCatalog(shared_catalog, Identifier("db", "tbl")) + .SetReadFieldNames({"name", "phone"}) + .SetPredicate(predicate); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context, scan_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, scan->CreatePlan()); + + // A worker without the catalog tells the read that the table enables query auth. + ReadContextBuilder read_builder(table_path); + read_builder.AddOption(Options::QUERY_AUTH_ENABLED, "true") + .SetReadFieldNames({"name", "phone"}) + .SetPredicate(predicate) + .EnablePredicateFilter(true); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context, read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, read->CreateReader(plan->Splits())); + +The read does not take the branch from the splits or from +``query-auth.enabled``: a worker reading a branch names it itself, as the scan +did. As for a scan, a branch named by the ``branch`` option or ``WithBranch`` is +carried into the identifier of a read built ``WithCatalog``, so the catalog +hands out the credentials of that branch. + +.. code-block:: cpp + + // Plan branch "dev"; Identifier("db", "tbl", "dev") names it as well. + ScanContextBuilder dev_scan_builder(table_path); + dev_scan_builder.WithCatalog(shared_catalog, Identifier("db", "tbl")) + .AddOption(Options::BRANCH, "dev") + .SetReadFieldNames({"name", "phone"}); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> dev_scan_context, + dev_scan_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> dev_scan, + TableScan::Create(std::move(dev_scan_context))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> dev_plan, dev_scan->CreatePlan()); + + // The worker reads the same branch. + ReadContextBuilder dev_read_builder(table_path); + dev_read_builder.AddOption(Options::QUERY_AUTH_ENABLED, "true") + .AddOption(Options::BRANCH, "dev") + .SetReadFieldNames({"name", "phone"}); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> dev_read_context, + dev_read_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> dev_read, + TableRead::Create(std::move(dev_read_context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> dev_reader, + dev_read->CreateReader(dev_plan->Splits())); The C++ REST catalog covers the database, table, snapshot and commit operations -of the ``Catalog`` API. The parts of the Java REST catalog that have no C++ -counterpart yet — altering a database or a table, views, functions, partitions, -tags, branch management and consumers — are not supported. Creating, deleting -and merging a branch is what branch management covers; committing to a branch -that already exists is supported, as described below. +of the ``Catalog`` API, and the query auth of a table. The parts of the Java REST +catalog that have no C++ counterpart yet — altering a database or a table, views, +functions, partitions, tags, branch management and consumers — are not supported. +Creating, deleting and merging a branch is what branch management covers; +committing to a branch that already exists is supported, as described below. Committing to a branch ~~~~~~~~~~~~~~~~~~~~~~ @@ -380,7 +631,8 @@ for the table while the files were written for the branch. Names a catalog cannot tell apart from another object are rejected where a -builder sees them: a name differing from ``main`` by case alone, such as +builder sees them, by the builders of a scan and a read through ``WithCatalog`` +as well: a name differing from ``main`` by case alone, such as ``tbl$branch_MAIN`` or ``branch=MAIN``, which a catalog reads as the main branch while its snapshots go to a directory of their own; and a name holding a ``$``, such as ``branch=dev$options``, which ``tbl$branch_dev$options`` reads back as @@ -400,13 +652,17 @@ above. A write or a commit aimed at a branch reads that schema rather than the -catalog's, just as a read of that branch does, and asks the catalog for no schema -at all, so a catalog serving only the table's own current schema still takes the -branch's snapshots. The new snapshot records the schema id published there as -well, which is the id a read of the branch resolves under -``branch/branch-<name>``; only on the main branch is the recorded id the one the -catalog reports, as its current schema may live nowhere else. Data files record -the schema they were written with, as they do on the main branch. +catalog's and asks the catalog for no schema at all, so a catalog serving only +the table's own current schema still takes the branch's snapshots. A scan or +read also uses the fields of the schema published under the branch, but through +``WithCatalog`` it asks for the catalog's branch schema to check +``query-auth.enabled`` (see `Reading through the catalog`_). + +The new snapshot records the schema id published under the branch, which is the +id a read of the branch resolves under ``branch/branch-<name>``; only on the main +branch is the recorded id the one the catalog reports, as its current schema may +live nowhere else. Data files record the schema they were written with, as they +do on the main branch. Addressing a branch without a catalog client ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/include/paimon/defs.h b/include/paimon/defs.h index 88af98c..dbb1203 100644 --- a/include/paimon/defs.h +++ b/include/paimon/defs.h
@@ -324,6 +324,12 @@ /// "rest.commit.lock-enabled" on the table. Default value is "false". static const char REST_COMMIT_ACQUIRE_LOCK[]; + /// "query-auth.enabled" - Whether the catalog enforces row filters and column masks on reads + /// of the table. A scan of such a table asks the catalog built into its context for the rules + /// and attaches them to its splits, and a read applies them; a read path that would not apply + /// them is refused. The options of a read cannot disable it. Default value is "false". + static const char QUERY_AUTH_ENABLED[]; + /// "compaction.max-size-amplification-percent" - The size amplification is defined as the /// amount (in percentage) of additional storage needed to store a single byte of data in the /// merge tree for changelog mode table. Default value is 200.
diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 4ca68dd..95aafb8 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h
@@ -408,6 +408,8 @@ /// Whether to perform precise filtering according to predicates for data read from format /// reader. + /// @note On a table with `query-auth.enabled=true` the predicate is evaluated on the masked + /// values, and its conjuncts on masked columns are evaluated even when this is false. /// @param enabled Whether to enable precise filtering (default: false) /// @return Reference to this builder for method chaining. ReadContextBuilder& EnablePredicateFilter(bool enabled); @@ -566,11 +568,20 @@ /// @note If not set, use default file system (configured in `Options::FILE_SYSTEM`) ReadContextBuilder& WithFileSystem(const std::shared_ptr<FileSystem>& file_system); - /// Reads a native table through its own file system - including the per-table temporary - /// credentials a catalog that issues them hands out through `Catalog::GetTableFileSystem`. - /// This is a shorthand for `WithFileSystem(catalog->GetTableFileSystem(identifier))`; an - /// explicit `WithFileSystem()` takes precedence, so the catalog is not asked. - /// @param catalog Non-null catalog, read when `Finish()` builds the context. + /// Reads a native table through its catalog, which supplies its file system: + /// `catalog->GetTableFileSystem(identifier)`, with any per-table temporary credentials; an + /// explicit `WithFileSystem()` takes precedence, so the catalog is not asked for one. + /// `TableRead::Create()` reads `query-auth.enabled` from the schema the catalog holds for the + /// branch read. A branch named by a system table path is resolved before requesting the table's + /// file system and must agree with the identifier, `branch` option and `WithBranch()`. + /// The fields the read reads with still come from `SetTableSchema()` on the main + /// branch or from the latest schema under the table path. When the catalog does not enable it, + /// the read path must be `catalog->GetTableLocation()` of the branch read. + /// A `sys` identifier requires a global system table path. A branch the catalog + /// would read as another object, such as `MAIN` or one holding a `$`, is rejected before the + /// catalog is asked for anything. + /// @param catalog Non-null catalog, read when `Finish()` builds the context and when + /// `TableRead::Create()` runs. /// @param identifier The native table to read. /// @return Reference to this builder for method chaining. ReadContextBuilder& WithCatalog(const std::shared_ptr<Catalog>& catalog,
diff --git a/include/paimon/scan_context.h b/include/paimon/scan_context.h index efee527..93f3ba2 100644 --- a/include/paimon/scan_context.h +++ b/include/paimon/scan_context.h
@@ -62,6 +62,22 @@ const std::shared_ptr<Catalog>& catalog, const std::optional<Identifier>& identifier); + /// Also takes the columns the query reads, see `ScanContextBuilder::SetReadFieldNames()`. + ScanContext(const std::string& path, bool is_streaming_mode, std::optional<int32_t> limit, + const std::shared_ptr<ScanFilter>& scan_filter, + const std::shared_ptr<GlobalIndexResult>& global_index_result, + const std::shared_ptr<RealtimeContext>& realtime_context, + const std::shared_ptr<MemoryPool>& memory_pool, + const std::shared_ptr<Executor>& executor, + const std::shared_ptr<FileSystem>& specific_file_system, + const std::optional<std::string>& table_schema, + const std::map<std::string, std::string>& options, + const std::shared_ptr<Cache>& cache, + const std::shared_ptr<FormatTable>& format_table, + const std::shared_ptr<Catalog>& catalog, + const std::optional<Identifier>& identifier, + const std::optional<std::vector<std::string>>& read_field_names); + ~ScanContext(); const std::string& GetPath() const { @@ -127,6 +143,10 @@ return identifier_; } + const std::optional<std::vector<std::string>>& GetReadFieldNames() const { + return read_field_names_; + } + private: std::string path_; bool is_streaming_mode_; @@ -143,6 +163,7 @@ std::shared_ptr<FormatTable> format_table_; std::shared_ptr<Catalog> catalog_; std::optional<Identifier> identifier_; + std::optional<std::vector<std::string>> read_field_names_; }; /// Filter configuration for table scan operations @@ -197,6 +218,11 @@ const std::vector<std::map<std::string, std::string>>& partition_filters); /// Set a predicate for filtering data. ScanContextBuilder& SetPredicate(const std::shared_ptr<Predicate>& predicate); + /// Set the columns the query reads. A table with `query-auth.enabled=true` asks its catalog + /// for the rules of these columns plus those of the predicate and the partition filter; by + /// default it asks for all columns. An empty list asks for no other column, unlike the empty + /// list of `ReadContextBuilder::SetReadFieldNames()`, which reads every column. + ScanContextBuilder& SetReadFieldNames(const std::vector<std::string>& read_field_names); /// Sets the result of a global index search (e.g., row ids (may with scores) from a distributed /// index lookup). This is used to push down index-filtered row ids into the scan for efficient @@ -241,11 +267,23 @@ /// @note If not set, use default file system (configured in `Options::FILE_SYSTEM`) ScanContextBuilder& WithFileSystem(const std::shared_ptr<FileSystem>& file_system); - /// Plans a native table through its own file system - including the per-table temporary - /// credentials a catalog that issues them hands out through `Catalog::GetTableFileSystem`. - /// This is a shorthand for `WithFileSystem(catalog->GetTableFileSystem(identifier))`; an - /// explicit `WithFileSystem()` takes precedence, so the catalog is not asked. - /// @param catalog Non-null catalog, read when `Finish()` builds the context. + /// Plans a native table through its catalog, which supplies its file system and, for a table + /// with `query-auth.enabled`, its rules. The file system is + /// `catalog->GetTableFileSystem(identifier)`, with any per-table temporary credentials; an + /// explicit `WithFileSystem()` takes precedence, so the catalog is not asked for one. + /// `TableScan::Create()` reads `query-auth.enabled` from the schema the catalog holds for the + /// branch scanned; the fields the scan plans with still come from `SetTableSchema()` on the + /// main branch or from the latest schema under the table path. The scan path must be + /// `catalog->GetTableLocation()` of the branch scanned, so the catalog's setting is taken only + /// for the table it keeps there. A `sys` identifier requires a global system table path. + /// A query-auth scan checks the location again on every plan, along with whether its selected + /// and rule columns keep their field ids in the catalog's current schema. + /// A branch the `branch` option or a system table path names is carried into the identifier, + /// as `tbl$branch_dev`, before requesting the table's file system; an + /// identifier naming another branch, or a branch the catalog would read as another object, such + /// as `MAIN` or one holding a `$`, is rejected before the catalog is asked for anything. + /// @param catalog Non-null catalog, read when `Finish()` builds the context and when + /// `TableScan::Create()` runs. /// @param identifier The native table to scan. /// @return Reference to this builder for method chaining. ScanContextBuilder& WithCatalog(const std::shared_ptr<Catalog>& catalog,
diff --git a/include/paimon/table/source/split.h b/include/paimon/table/source/split.h index 7f296fb..bcf535a 100644 --- a/include/paimon/table/source/split.h +++ b/include/paimon/table/source/split.h
@@ -39,7 +39,8 @@ /// This split can be a `DataSplit` (for direct data file reads), an `IndexedSplit` (for reads /// leveraging global indexes), or a `FormatDataSplit` (for a format table's plain data files). /// Only the first two are serializable; a `FormatDataSplit` is in-memory only and `Serialize()` -/// refuses it. +/// refuses it. A split of a table with `query-auth.enabled=true` wraps one of the first two with +/// the row filters and column masks it is read with. class PAIMON_EXPORT Split { public: virtual ~Split() = default;
diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 72866e3..c89d7f4 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt
@@ -269,6 +269,7 @@ core/catalog/catalog_utils.cpp core/catalog/file_system_catalog.cpp core/catalog/identifier.cpp + core/catalog/table_query_auth_result.cpp core/core_options.cpp core/deletionvectors/deletion_vector.cpp core/deletionvectors/deletion_file_writer.cpp @@ -485,6 +486,8 @@ core/table/source/merge_tree_split_generator.cpp core/table/source/data_evolution_split_generator.cpp core/table/source/plan_impl.cpp + core/table/source/query_auth_table_read.cpp + core/table/source/query_auth_table_scan.cpp core/table/source/realtime_table_scan.cpp core/table/source/snapshot/snapshot_reader.cpp core/table/source/startup_mode.cpp @@ -852,6 +855,7 @@ core/catalog/file_system_catalog_test.cpp core/catalog/catalog_test.cpp core/catalog/identifier_test.cpp + core/catalog/table_query_auth_result_test.cpp core/compact/compact_deletion_file_test.cpp core/compact/compact_future_manager_test.cpp core/compact/noop_compact_manager_test.cpp @@ -1046,6 +1050,8 @@ core/table/sink/commit_message_impl_test.cpp core/table/source/fallback_data_split_test.cpp core/table/source/primary_key_sorted_index_scan_test.cpp + core/table/source/query_auth_split_test.cpp + core/table/source/query_auth_table_read_test.cpp core/table/source/table_read_test.cpp core/table/source/append_count_reader_test.cpp core/table/source/pk_count_reader_test.cpp
diff --git a/src/paimon/common/defs.cpp b/src/paimon/common/defs.cpp index dc88b6b..b05f568 100644 --- a/src/paimon/common/defs.cpp +++ b/src/paimon/common/defs.cpp
@@ -91,6 +91,7 @@ const char Options::COMMIT_DISCARD_DUPLICATE_FILES[] = "commit.discard-duplicate-files"; const char Options::REST_COMMIT_LOCK_ENABLED[] = "rest.commit.lock-enabled"; const char Options::REST_COMMIT_ACQUIRE_LOCK[] = "rest.commit.acquire-lock"; +const char Options::QUERY_AUTH_ENABLED[] = "query-auth.enabled"; const char Options::SEQUENCE_FIELD[] = "sequence.field"; const char Options::SEQUENCE_FIELD_SORT_ORDER[] = "sequence.field.sort-order"; const char Options::MERGE_ENGINE[] = "merge-engine";
diff --git a/src/paimon/common/predicate/like.cpp b/src/paimon/common/predicate/like.cpp index d6ee817..b935ae9 100644 --- a/src/paimon/common/predicate/like.cpp +++ b/src/paimon/common/predicate/like.cpp
@@ -44,11 +44,6 @@ return 1; // invalid continuation byte, treat as single byte } -inline bool IsJavaRegexLineTerminator(const std::string& code_point) { - return code_point == "\n" || code_point == "\r" || code_point == "\xC2\x85" || - code_point == "\xE2\x80\xA8" || code_point == "\xE2\x80\xA9"; -} - } // namespace Result<bool> Like::MatchString(const std::string& field, const std::string& pattern) { @@ -153,7 +148,8 @@ if (wild && pc == "%") { dp[j] = dp[j - 1] || dp[j]; } else if (wild && pc == "_") { - dp[j] = prev && !IsJavaRegexLineTerminator(field_char); + // As Java `Like`, which matches `_` with `(?s:.)`: any one code point. + dp[j] = prev; } else { dp[j] = (pc == field_char) ? prev : false; }
diff --git a/src/paimon/common/predicate/literal_converter.cpp b/src/paimon/common/predicate/literal_converter.cpp index 2b1a1eb..4604bdd 100644 --- a/src/paimon/common/predicate/literal_converter.cpp +++ b/src/paimon/common/predicate/literal_converter.cpp
@@ -33,6 +33,7 @@ #include "arrow/type.h" #include "arrow/type_traits.h" #include "arrow/util/decimal.h" +#include "arrow/util/int_util_overflow.h" #include "fmt/format.h" #include "paimon/common/data/binary_string.h" #include "paimon/common/data/internal_row.h" @@ -379,6 +380,48 @@ return array; } +// The number of `unit`s since the epoch `timestamp` names, or an error when that number does not +// fit int64. Epoch nanoseconds only reach the years 1677 to 2262, so a timestamp outside them +// still has milliseconds but no nanoseconds to write. +Result<int64_t> TimestampToCheckedInteger(const Timestamp& timestamp, arrow::TimeUnit::type unit) { + int64_t millis = timestamp.GetMillisecond(); + const int64_t nanos = timestamp.GetNanoOfMillisecond(); + int64_t factor = 1; + int64_t remainder = 0; + switch (unit) { + case arrow::TimeUnit::MILLI: + return millis; + case arrow::TimeUnit::MICRO: + factor = 1000; + remainder = nanos / 1000; + break; + case arrow::TimeUnit::NANO: + factor = 1000000; + remainder = nanos; + break; + default: + return Status::Invalid(fmt::format("Not support writing a timestamp literal in {}", + arrow::timestamp(unit)->ToString())); + } + // The earliest values that fit have a negative millisecond whose product alone does not, so a + // negative millisecond lends one `factor` to the remainder first. + if (millis < 0 && remainder > 0) { + ++millis; + remainder -= factor; + } + int64_t scaled = 0; + int64_t value = 0; + if (arrow::internal::MultiplyWithOverflow(millis, factor, &scaled) || + arrow::internal::AddWithOverflow(scaled, remainder, &value)) { + // `Timestamp::ToString` cannot format an instant this far from the epoch. + return Status::Invalid(fmt::format( + "Timestamp literal of {} milliseconds and {} nanoseconds since the epoch does not fit " + "int64 in {}", + timestamp.GetMillisecond(), nanos, arrow::timestamp(unit)->ToString())); + } + return value; +} + // Writes the literals to a timestamp array of the finest time unit that keeps every value, one // `is_in` can only compare against a column of that very unit. One array holds one arrow type, so // a non null literal is needed to settle the unit. @@ -396,18 +439,17 @@ "Not support converting literals of TIMESTAMP type to an arrow array without a non " "null literal to take the time unit from"); } - arrow::TimestampBuilder builder( - arrow::timestamp(LiteralConverter::MinRequiredTimeUnit(literals)), pool); - const DateTimeUtils::TimeType time_type = DateTimeUtils::GetTimeTypeFromArrowType( - checked_pointer_cast<arrow::TimestampType>(builder.type())); + const arrow::TimeUnit::type unit = LiteralConverter::MinRequiredTimeUnit(literals); + arrow::TimestampBuilder builder(arrow::timestamp(unit), pool); PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast<int64_t>(literals.size()))); for (const auto& literal : literals) { if (literal.IsNull()) { PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); continue; } - PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append( - DateTimeUtils::TimestampToInteger(literal.GetValue<Timestamp>(), time_type))); + PAIMON_ASSIGN_OR_RAISE(int64_t value, + TimestampToCheckedInteger(literal.GetValue<Timestamp>(), unit)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(value)); } std::shared_ptr<arrow::Array> array; PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array));
diff --git a/src/paimon/common/predicate/literal_converter.h b/src/paimon/common/predicate/literal_converter.h index 9692d5a..c393855 100644 --- a/src/paimon/common/predicate/literal_converter.h +++ b/src/paimon/common/predicate/literal_converter.h
@@ -63,7 +63,8 @@ /// precision and the scale its literals carry, so at least one literal has to be non /// null and every non null one has to carry the same pair. `TIMESTAMP` is written with /// the finest time unit that keeps every value, so at least one literal has to be non - /// null to settle the unit. + /// null to settle the unit, and a literal whose value in that unit does not fit int64 + /// fails. static Result<std::shared_ptr<arrow::Array>> ConvertLiteralsToArray( const FieldType& field_type, const std::vector<Literal>& literals, arrow::MemoryPool* pool);
diff --git a/src/paimon/common/predicate/literal_converter_test.cpp b/src/paimon/common/predicate/literal_converter_test.cpp index 201f2ae..c862960 100644 --- a/src/paimon/common/predicate/literal_converter_test.cpp +++ b/src/paimon/common/predicate/literal_converter_test.cpp
@@ -21,6 +21,7 @@ #include <algorithm> #include <cmath> #include <cstddef> +#include <limits> #include <variant> #include "arrow/api.h" @@ -597,6 +598,38 @@ arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"([null, null])").ValueOrDie()); } +TEST_F(LiteralConverterTest, TestLiteralsToArrayRejectsTimestampsOutOfRange) { + // INT64_MAX and INT64_MIN epoch nanoseconds + CheckLiteralsToArray( + FieldType::TIMESTAMP, + {Literal(Timestamp(9223372036854l, 775807)), Literal(Timestamp(-9223372036855l, 224192))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::NANO), + R"([9223372036854775807, -9223372036854775808])") + .ValueOrDie()); + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray( + FieldType::TIMESTAMP, {Literal(Timestamp(9223372036854l, 775808))}, + arrow::default_memory_pool()), + "does not fit int64"); + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray( + FieldType::TIMESTAMP, {Literal(Timestamp(-9223372036855l, 224191))}, + arrow::default_memory_pool()), + "does not fit int64"); + // 2500-01-01T00:00:00.000000001, beyond epoch nanoseconds + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray( + FieldType::TIMESTAMP, {Literal(Timestamp(16725225600000l, 1))}, + arrow::default_memory_pool()), + "does not fit int64"); + CheckLiteralsToArray(FieldType::TIMESTAMP, {Literal(Timestamp(16725225600000l, 1000))}, + arrow::ipc::internal::json::ArrayFromJSON( + arrow::timestamp(arrow::TimeUnit::MICRO), R"([16725225600000001])") + .ValueOrDie()); + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray( + FieldType::TIMESTAMP, {Literal(Timestamp(std::numeric_limits<int64_t>::max(), 1000))}, + arrow::default_memory_pool()), + "does not fit int64"); +} + TEST_F(LiteralConverterTest, TestLiteralsToArrayUnsupportedType) { // a timestamp takes its arrow type from a literal, so it needs one that is not null ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::TIMESTAMP, {},
diff --git a/src/paimon/common/predicate/multi_literals_leaf_function.cpp b/src/paimon/common/predicate/multi_literals_leaf_function.cpp index d423e57..a2e6d58 100644 --- a/src/paimon/common/predicate/multi_literals_leaf_function.cpp +++ b/src/paimon/common/predicate/multi_literals_leaf_function.cpp
@@ -168,7 +168,8 @@ Result<std::shared_ptr<arrow::Array>> value_set = LiteralConverter::ConvertLiteralsToArray(field_type, literals, pool); // A failure only says the value set is not there, and the row by row path still is. Decimal - // literals of mixed precision and scale end up here, anything else is an arrow failure. + // literals of mixed precision and scale and a timestamp literal that does not fit int64 in the + // unit end up here, anything else is an arrow failure. if (!value_set.ok()) { return nullptr; }
diff --git a/src/paimon/common/predicate/null_false_leaf_binary_function.cpp b/src/paimon/common/predicate/null_false_leaf_binary_function.cpp index 7aa9575..b4db667 100644 --- a/src/paimon/common/predicate/null_false_leaf_binary_function.cpp +++ b/src/paimon/common/predicate/null_false_leaf_binary_function.cpp
@@ -154,8 +154,9 @@ /// settles whether a decimal or a timestamp scalar can be compared against it. /// @param pool The pool the scalar is allocated from. /// @return `nullptr` when the literal cannot be compared by a kernel, which covers a decimal column -/// of another scale and a timestamp column of another unit or with a time zone. This never -/// fails. +/// of another scale, a timestamp column of another unit or with a time zone, and a +/// timestamp literal that does not fit int64 in the time unit it is written with. This +/// never fails. std::shared_ptr<arrow::Scalar> MakeComparisonScalar(const Literal& literal, FieldType field_type, const arrow::DataType& column_type, arrow::MemoryPool* pool) {
diff --git a/src/paimon/common/predicate/predicate_test.cpp b/src/paimon/common/predicate/predicate_test.cpp index 97fef15..8094668 100644 --- a/src/paimon/common/predicate/predicate_test.cpp +++ b/src/paimon/common/predicate/predicate_test.cpp
@@ -1765,26 +1765,30 @@ ASSERT_FALSE(predicate->Test(arrow_schema, CreateStringRow({"helloworld"})).value()); } -TEST_F(PredicateTest, TestLikeJavaRegexLineTerminatorSemantics) { +TEST_F(PredicateTest, TestLikeWildcardsMatchLineTerminators) { auto arrow_schema = arrow::schema(arrow::FieldVector({arrow::field("f0", arrow::utf8())})); - - // Java regex '.' does not match line terminators, so '_' should not match them either. - ASSERT_OK_AND_ASSIGN(auto predicate_base, - PredicateBuilder::Like( - /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, - Literal(FieldType::STRING, "_", 1))); - auto predicate = std::dynamic_pointer_cast<PredicateFilter>(predicate_base); - ASSERT_FALSE(predicate->Test(arrow_schema, CreateStringRow({"\n"})).value()); - ASSERT_FALSE(predicate->Test(arrow_schema, CreateStringRow({"\r"})).value()); - - // Java LIKE '%' uses (?s:.*), so it should still match line terminators. - ASSERT_OK_AND_ASSIGN(predicate_base, - PredicateBuilder::Like( - /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, - Literal(FieldType::STRING, "%", 1))); - predicate = std::dynamic_pointer_cast<PredicateFilter>(predicate_base); - ASSERT_TRUE(predicate->Test(arrow_schema, CreateStringRow({"\n"})).value()); - ASSERT_TRUE(predicate->Test(arrow_schema, CreateStringRow({"\r"})).value()); + auto like = [&](const std::string& pattern, const std::string& value) -> Result<bool> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> predicate, + PredicateBuilder::Like( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, pattern.data(), pattern.size()))); + auto filter = std::dynamic_pointer_cast<PredicateFilter>(predicate); + if (filter == nullptr) { + return Status::Invalid("the LIKE predicate is not a filter"); + } + return filter->Test(arrow_schema, CreateStringRow({value})); + }; + for (const std::string line_terminator : + {"\n", "\r", "\xC2\x85", "\xE2\x80\xA8", "\xE2\x80\xA9"}) { + ASSERT_OK_AND_ASSIGN(bool matched, like("a_b", "a" + line_terminator + "b")); + ASSERT_TRUE(matched) << line_terminator; + ASSERT_OK_AND_ASSIGN(matched, like("%", line_terminator)); + ASSERT_TRUE(matched) << line_terminator; + } + ASSERT_OK_AND_ASSIGN(bool matched, like("a_b", "a\r\nb")); + ASSERT_FALSE(matched); + ASSERT_OK_AND_ASSIGN(matched, like("a__b", "a\r\nb")); + ASSERT_TRUE(matched); } TEST_F(PredicateTest, TestCompound) {
diff --git a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp index a27ccd5..ec0bc26 100644 --- a/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp +++ b/src/paimon/common/reader/blob_fallback_batch_reader_test.cpp
@@ -324,6 +324,7 @@ // This mock ignores the context and returns two rows for a one-row read/skip request. auto oversized_reader = std::make_unique<MockFileBatchReader>(MakeBlobStruct({"b0", "b1"}), struct_type_, 1024); + oversized_reader->EnableRandomizeBatchSize(false); std::vector<std::vector<BlobFallbackBatchReader::Segment>> groups; groups.push_back(MakeGroup( {SegmentSpec::File(covered ? BlobRows{"u0", "u1"} : BlobRows{"PH", "PH"})}, 1024));
diff --git a/src/paimon/common/types/data_type_json_parser.cpp b/src/paimon/common/types/data_type_json_parser.cpp index 1416ae3..c05110f 100644 --- a/src/paimon/common/types/data_type_json_parser.cpp +++ b/src/paimon/common/types/data_type_json_parser.cpp
@@ -29,12 +29,14 @@ #include <utility> #include <vector> +#include "arrow/type.h" #include "fmt/format.h" #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/types/data_field.h" #include "paimon/common/types/data_type.h" #include "paimon/common/types/vector_type.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/string_utils.h" #include "paimon/data/decimal.h" @@ -575,7 +577,15 @@ } PAIMON_RETURN_NOT_OK(NextToken(TokenType::END_PARAMETER)); } - return arrow::decimal128(precision, scale); + // `arrow::decimal128` aborts the process on a precision out of range, `Make` reports it. + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::DataType> type, + arrow::Decimal128Type::Make(precision, scale)); + if (scale < 0 || scale > precision) { + return Status::Invalid(fmt::format( + "Decimal scale must be between 0 and the precision {} (both inclusive), but was {}", + precision, scale)); + } + return type; } Result<std::shared_ptr<arrow::DataType>> TokenParser::ParseDoubleType() { @@ -757,8 +767,9 @@ const std::string& name, const rapidjson::Value& type_json_value) { bool nullable = true; AtomicTypeAttributes attributes; + const std::string type_str(type_json_value.GetString(), type_json_value.GetStringLength()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::DataType> type, - ParseAtomicType(type_json_value.GetString(), &nullable, &attributes)); + ParseAtomicType(type_str, &nullable, &attributes)); if (attributes.is_blob) { return BlobUtils::ToArrowField(name, nullable); } else if (attributes.is_variant) { @@ -776,11 +787,15 @@ Result<std::shared_ptr<arrow::Field>> DataTypeJsonParser::ParseComplexTypeField( const std::string& name, const rapidjson::Value& type_json_value, FieldIdAssigner* field_id_assigner) { - if (!type_json_value.HasMember("type")) { - return Status::Invalid("complex data type must have type"); + if (!type_json_value.HasMember("type") || !type_json_value["type"].IsString()) { + return Status::Invalid("complex data type must have a string type"); } - std::string type_str = type_json_value["type"].GetString(); + const auto& type_value = type_json_value["type"]; + const std::string type_str(type_value.GetString(), type_value.GetStringLength()); + if (type_str.find('\0') != std::string::npos) { + return Status::Invalid("complex data type cannot contain a NUL byte"); + } bool nullable = true; if (type_str.find("NOT NULL") != std::string::npos) { nullable = false;
diff --git a/src/paimon/common/types/data_type_json_parser_test.cpp b/src/paimon/common/types/data_type_json_parser_test.cpp index 7c70289..9c11956 100644 --- a/src/paimon/common/types/data_type_json_parser_test.cpp +++ b/src/paimon/common/types/data_type_json_parser_test.cpp
@@ -228,8 +228,21 @@ } } +TEST(DataTypeJsonParserTest, ParseTypeRejectsEmbeddedNul) { + for (const char* json : + {R"("INT\u0000suffix")", R"({"type": "ARRAY\u0000suffix", "element": "INT"})", + R"({"type": "ARRAY", "element": "STRING\u0000suffix"})"}) { + rapidjson::Document doc; + doc.Parse(json); + ASSERT_FALSE(doc.HasParseError()); + ASSERT_NOK(DataTypeJsonParser::ParseType("field", doc)); + } +} + TEST(DataTypeJsonParserTest, ParseTypeComplexTypeFailure) { const std::vector<std::pair<const char*, const char*>> test_cases = { + {R"({"type":123})", "complex data type must have a string type"}, + {R"({"element":"INT"})", "complex data type must have a string type"}, {R"({"type":"ROW"})", "row data type must have fields"}, {R"({"type":"ROW","fields":{}})", "row data type must have fields"}, {R"({"type":"ROW","fields":[null]})", "data field must be an object"}, @@ -287,6 +300,8 @@ {"NUMERIC", arrow::decimal128(10, 0)}, {"NUMERIC(10)", arrow::decimal128(10, 0)}, {"NUMERIC(10, 3)", arrow::decimal128(10, 3)}, + {"DECIMAL(1, 0)", arrow::decimal128(1, 0)}, + {"DECIMAL(38, 38)", arrow::decimal128(38, 38)}, {"TIMESTAMP(0)", arrow::timestamp(arrow::TimeUnit::SECOND)}, {"TIMESTAMP(3)", arrow::timestamp(arrow::TimeUnit::MILLI)}, {"TIMESTAMP(6)", arrow::timestamp(arrow::TimeUnit::MICRO)}, @@ -344,6 +359,18 @@ ASSERT_NOK_WITH_MSG(DataTypeJsonParser::ParseType("field_name", value), "length must be between 1 and 2147483647"); } + for (const char* invalid_type : {"DECIMAL(0)", "DECIMAL(39)", "NUMERIC(100, 2)"}) { + rapidjson::Document invalid_doc; + rapidjson::Value value(invalid_type, invalid_doc.GetAllocator()); + ASSERT_NOK_WITH_MSG(DataTypeJsonParser::ParseType("field_name", value), + "Decimal precision out of range"); + } + for (const char* invalid_type : {"DECIMAL(10, 11)", "DECIMAL(38, 39)"}) { + rapidjson::Document invalid_doc; + rapidjson::Value value(invalid_type, invalid_doc.GetAllocator()); + ASSERT_NOK_WITH_MSG(DataTypeJsonParser::ParseType("field_name", value), + "Decimal scale must be between 0 and the precision"); + } { rapidjson::Document invalid_doc; rapidjson::Value value("TIMESTAMP(4)", invalid_doc.GetAllocator());
diff --git a/src/paimon/common/utils/rapidjson_util.h b/src/paimon/common/utils/rapidjson_util.h index d413f25..bac1b9b 100644 --- a/src/paimon/common/utils/rapidjson_util.h +++ b/src/paimon/common/utils/rapidjson_util.h
@@ -23,8 +23,10 @@ #include <cstdint> #include <map> #include <optional> +#include <set> #include <stdexcept> #include <string> +#include <string_view> #include "paimon/common/utils/string_utils.h" #include "paimon/status.h" @@ -76,6 +78,10 @@ if (!obj) { return Status::Invalid("deserialize failed: obj is nullptr"); } + // RapidJSON treats an unescaped NUL as end-of-input, even with an explicit length. + if (json_str.find('\0') != std::string::npos) { + return Status::Invalid("deserialize failed: JSON contains an unescaped NUL byte"); + } if constexpr (std::is_same_v<T, std::map<std::string, std::string>>) { PAIMON_ASSIGN_OR_RAISE(*obj, MapFromJsonString(json_str)); } else { @@ -115,6 +121,10 @@ template <typename T> static T DeserializeValue(const rapidjson::Value& value); + // Fails when an object in `value`, at any depth, names a member twice: RapidJSON keeps both + // and a lookup finds the first. + static Status CheckNoDuplicateMembers(const rapidjson::Value& value); + private: static inline bool ToJson(const rapidjson::Value& value, std::string* json_str) { assert(json_str); @@ -189,6 +199,25 @@ } }; +inline Status RapidJsonUtil::CheckNoDuplicateMembers(const rapidjson::Value& value) { + if (value.IsArray()) { + for (const rapidjson::Value& item : value.GetArray()) { + PAIMON_RETURN_NOT_OK(CheckNoDuplicateMembers(item)); + } + } else if (value.IsObject()) { + std::set<std::string_view> names; + for (auto it = value.MemberBegin(); it != value.MemberEnd(); ++it) { + std::string_view name(it->name.GetString(), it->name.GetStringLength()); + if (!names.insert(name).second) { + return Status::Invalid("JSON object names member '", std::string(name), + "' more than once"); + } + PAIMON_RETURN_NOT_OK(CheckNoDuplicateMembers(it->value)); + } + } + return Status::OK(); +} + template <typename T> inline rapidjson::Value RapidJsonUtil::SerializeValue( const T& obj, rapidjson::Document::AllocatorType* allocator) { @@ -202,7 +231,8 @@ } } else { if constexpr (std::is_same_v<T, std::string>) { - return rapidjson::Value(obj.c_str(), *allocator); + return rapidjson::Value(obj.data(), static_cast<rapidjson::SizeType>(obj.size()), + *allocator); } else if constexpr (std::is_arithmetic_v<T>) { return rapidjson::Value(obj); } else if constexpr (is_map<T>::value) {
diff --git a/src/paimon/common/utils/rapidjson_util_test.cpp b/src/paimon/common/utils/rapidjson_util_test.cpp index c5011d9..2a27743 100644 --- a/src/paimon/common/utils/rapidjson_util_test.cpp +++ b/src/paimon/common/utils/rapidjson_util_test.cpp
@@ -144,4 +144,41 @@ ASSERT_EQ(m1, m2); } +TEST(RapidJsonUtilTest, TestRejectUnescapedNul) { + std::map<std::string, std::string> result; + for (const std::string& suffix : {std::string(1, '\0'), std::string(1, '\0') + "not JSON"}) { + ASSERT_NOK_WITH_MSG(RapidJsonUtil::FromJsonString("{}" + suffix, &result), + "JSON contains an unescaped NUL byte"); + } +} + +TEST(RapidJsonUtilTest, TestSerializeStringWithEmbeddedNul) { + rapidjson::Document doc; + const std::string text("name\0secret", 11); + rapidjson::Value value = RapidJsonUtil::SerializeValue(text, &doc.GetAllocator()); + ASSERT_EQ(text.size(), value.GetStringLength()); + ASSERT_EQ(text, RapidJsonUtil::DeserializeValue<std::string>(value)); + + const std::map<std::string, std::string> map = {{text, text}}; + rapidjson::Value object = RapidJsonUtil::SerializeValue(map, &doc.GetAllocator()); + ASSERT_EQ(map, (RapidJsonUtil::DeserializeValue<std::map<std::string, std::string>>(object))); +} + +TEST(RapidJsonUtilTest, TestCheckNoDuplicateMembers) { + auto check = [](const std::string& json) { + rapidjson::Document doc; + doc.Parse(json.c_str(), json.size()); + EXPECT_FALSE(doc.HasParseError()) << json; + return RapidJsonUtil::CheckNoDuplicateMembers(doc); + }; + ASSERT_OK(check(R"({"a": {"a": 1, "b": [{"a": 1}, {"a": 2}]}, "b": null})")); + ASSERT_OK(check(R"([1, "a", {"a": [], "A": {}}])")); + ASSERT_OK(check("{}")); + + ASSERT_NOK_WITH_MSG(check(R"({"a": [], "a": [1]})"), "JSON object names member 'a' more than"); + ASSERT_NOK_WITH_MSG(check(R"({"a": {"b": 1, "c": {"d": 1, "d": 1}}})"), "member 'd'"); + ASSERT_NOK_WITH_MSG(check(R"([{}, [{"e": 1, "f": 2, "e": 3}]])"), "member 'e'"); + ASSERT_NOK_WITH_MSG(check(R"({"filter": 1, "fil\u0074er": 2})"), "member 'filter'"); +} + } // namespace paimon::test
diff --git a/src/paimon/core/catalog/catalog_utils.cpp b/src/paimon/core/catalog/catalog_utils.cpp index 1f2a380..1d6a408 100644 --- a/src/paimon/core/catalog/catalog_utils.cpp +++ b/src/paimon/core/catalog/catalog_utils.cpp
@@ -26,6 +26,7 @@ #include "fmt/format.h" #include "paimon/catalog/catalog.h" #include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/string_utils.h" #include "paimon/core/options/table_type.h" #include "paimon/core/utils/branch_manager.h" #include "paimon/defs.h" @@ -42,6 +43,15 @@ action, identifier.ToString())); } +bool SameLocation(const Path& left, const Path& right) { + auto scheme = [](const Path& path) { + std::string lower = StringUtils::ToLowerCase(path.scheme); + return lower.empty() ? std::string("file") : lower; + }; + return scheme(left) == scheme(right) && left.authority == right.authority && + left.path == right.path; +} + } // namespace bool CatalogUtils::IsSystemDatabase(const std::string& db_name) { @@ -75,6 +85,53 @@ return Identifier(identifier.GetDatabaseName(), data_table_name, branch); } +Result<Identifier> CatalogUtils::ResolvedBranchIdentifier(const Identifier& identifier, + const std::string& branch) { + if (IsSystemDatabase(identifier.GetDatabaseName())) { + return identifier; + } + PAIMON_ASSIGN_OR_RAISE(std::optional<std::string> identifier_branch, + identifier.GetBranchName()); + if (identifier_branch) { + return identifier; + } + PAIMON_ASSIGN_OR_RAISE(Identifier branch_identifier, BranchIdentifier(identifier, branch)); + PAIMON_ASSIGN_OR_RAISE(std::optional<std::string> system_table, + identifier.GetSystemTableName()); + if (!system_table) { + return branch_identifier; + } + return Identifier(branch_identifier.GetDatabaseName(), + fmt::format("{}{}{}", branch_identifier.GetObjectName(), + Identifier::kSystemTableSplitter, system_table.value())); +} + +Result<bool> CatalogUtils::SameTableLocation(const std::string& left, const std::string& right) { + PAIMON_ASSIGN_OR_RAISE(Path left_path, PathUtil::ToPath(left)); + PAIMON_ASSIGN_OR_RAISE(Path right_path, PathUtil::ToPath(right)); + return SameLocation(left_path, right_path); +} + +Status CatalogUtils::CheckPathIsTableLocation(const Identifier& identifier, + const Result<std::string>& location, + const std::string& path, + const std::string& operation) { + if (!location.ok()) { + return location.status().WithMessage(fmt::format( + "Cannot check that {} is the location of table {}, whose query auth settings the {} " + "applies: {}", + path, identifier.GetFullName(), operation, location.status().message())); + } + PAIMON_ASSIGN_OR_RAISE(bool same_location, SameTableLocation(location.value(), path)); + if (!same_location) { + return Status::Invalid(fmt::format( + "The {} targets {}, but the catalog keeps table {}, whose query auth settings it would " + "apply, at {}. Use the location the catalog returns for the table.", + operation, path, identifier.GetFullName(), location.value())); + } + return Status::OK(); +} + Status CatalogUtils::CheckNotBranch(const Identifier& identifier, const std::string& action) { PAIMON_ASSIGN_OR_RAISE(std::optional<std::string> branch, identifier.GetBranchName()); if (branch) {
diff --git a/src/paimon/core/catalog/catalog_utils.h b/src/paimon/core/catalog/catalog_utils.h index 565a6ae..93c5653 100644 --- a/src/paimon/core/catalog/catalog_utils.h +++ b/src/paimon/core/catalog/catalog_utils.h
@@ -58,6 +58,25 @@ static Result<Identifier> BranchIdentifier(const Identifier& identifier, const std::string& branch); + /// Returns `identifier` naming `branch`, the branch resolved for an operation on it, so that + /// every catalog request of the operation goes to that branch: an identifier of the system + /// database or one already naming a branch is returned as it is, and any other keeps its + /// system table suffix, as `tbl$branch_dev$files`. + static Result<Identifier> ResolvedBranchIdentifier(const Identifier& identifier, + const std::string& branch); + + /// Compares scheme, authority and normalized path, treating a missing scheme as "file". + static Result<bool> SameTableLocation(const std::string& left, const std::string& right); + + /// Fails unless `path` is `location`, where the catalog keeps `identifier`, so the query auth + /// settings the catalog serves for one table are never applied to the files of another. + /// @param location The location the catalog returned for `identifier`, or the error it failed + /// with. + /// @param operation The operation reading `path`, as "scan", named in the error message. + static Status CheckPathIsTableLocation(const Identifier& identifier, + const Result<std::string>& location, + const std::string& path, const std::string& operation); + /// Fails when `db_name` cannot be used as a single path component, which is required to /// keep the database path under the warehouse. static Status CheckValidDatabaseName(const std::string& db_name);
diff --git a/src/paimon/core/catalog/query_auth_catalog.h b/src/paimon/core/catalog/query_auth_catalog.h new file mode 100644 index 0000000..e4a7652 --- /dev/null +++ b/src/paimon/core/catalog/query_auth_catalog.h
@@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <memory> +#include <optional> +#include <string> +#include <vector> + +#include "paimon/catalog/catalog.h" +#include "paimon/catalog/identifier.h" +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/result.h" + +namespace paimon { + +/// Serves the row filters and column masks of tables with `query-auth.enabled=true`, mirroring +/// Java `Catalog.authTableQuery`. +/// +/// Kept internal because `TableQueryAuthResult` is not an installed API type. +class QueryAuthCatalog { + public: + virtual ~QueryAuthCatalog() = default; + + /// Loads the rules one query of a table is subject to. + /// + /// @param identifier The table the query reads, a branch of it included. A system table is + /// refused. + /// @param select The columns the query reads, or null for all of them. + /// @return The parsed rules; a rule this library cannot apply fails the call. + virtual Result<std::shared_ptr<TableQueryAuthResult>> AuthTableQuery( + const Identifier& identifier, + const std::optional<std::vector<std::string>>& select) const = 0; + + /// Loads the schema of a table together with its location, so both describe the same table + /// even when it is recreated meanwhile. + /// + /// @param identifier The table, a branch of it included. + /// @param location Set to the location the catalog keeps the table at. + virtual Result<std::shared_ptr<Schema>> LoadTableSchemaAndLocation( + const Identifier& identifier, std::string* location) const = 0; +}; + +/// Returns the catalog as a `QueryAuthCatalog`, or null when it does not support the interface. +inline const QueryAuthCatalog* AsQueryAuthCatalog(const std::shared_ptr<Catalog>& catalog) { + return dynamic_cast<const QueryAuthCatalog*>(catalog.get()); +} + +} // namespace paimon
diff --git a/src/paimon/core/catalog/table_query_auth_result.cpp b/src/paimon/core/catalog/table_query_auth_result.cpp new file mode 100644 index 0000000..7806845 --- /dev/null +++ b/src/paimon/core/catalog/table_query_auth_result.cpp
@@ -0,0 +1,1369 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/catalog/table_query_auth_result.h" + +#include <algorithm> +#include <cmath> +#include <cstdint> +#include <deque> +#include <limits> +#include <optional> +#include <string_view> +#include <utility> + +#include "arrow/api.h" +#include "arrow/util/decimal.h" +#include "fmt/format.h" +#include "fmt/ranges.h" +#include "paimon/common/data/binary_string.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/types/data_type_json_parser.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/rapidjson_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/predicate_utils.h" +#include "rapidjson/document.h" + +namespace paimon { +namespace { + +using ReferencedTypes = std::map<std::string, std::shared_ptr<arrow::DataType>>; + +/// Fails when two rules read one column as different types. +Status RecordReferencedType(const std::string& name, const std::shared_ptr<arrow::DataType>& type, + ReferencedTypes* types) { + auto [iter, inserted] = types->emplace(name, type); + if (!inserted && !iter->second->Equals(*type)) { + return Status::Invalid(fmt::format("Query auth rules read column '{}' as both {} and {}.", + name, iter->second->ToString(), type->ToString())); + } + return Status::OK(); +} + +Result<std::string> GetString(const rapidjson::Value& object, const char* key, + const std::string& what) { + auto iter = object.FindMember(key); + if (iter == object.MemberEnd() || !iter->value.IsString()) { + return Status::Invalid(fmt::format("{} needs a string member '{}'.", what, key)); + } + return std::string(iter->value.GetString(), iter->value.GetStringLength()); +} + +/// A Java `FieldRef`: `{"index": 0, "name": "f0", "type": "INT"}`. +struct FieldRef { + int32_t index = 0; + std::string name; + std::shared_ptr<arrow::DataType> type; + /// `UNKNOWN` for a type a predicate cannot compare, such as TIME, which a mask still copies. + FieldType field_type = FieldType::UNKNOWN; +}; + +Result<FieldRef> ParseFieldRef(const rapidjson::Value& value, ReferencedTypes* types) { + if (!value.IsObject()) { + return Status::Invalid("A field reference of a query auth rule must be a JSON object."); + } + FieldRef ref; + PAIMON_ASSIGN_OR_RAISE(ref.name, GetString(value, "name", "A field reference")); + if (ref.name.empty()) { + return Status::Invalid("A field reference of a query auth rule cannot have an empty name."); + } + auto index_iter = value.FindMember("index"); + if (index_iter != value.MemberEnd() && index_iter->value.IsInt()) { + ref.index = index_iter->value.GetInt(); + } + auto type_iter = value.FindMember("type"); + if (type_iter == value.MemberEnd() || type_iter->value.IsNull()) { + return Status::Invalid( + fmt::format("The field reference to '{}' of a query auth rule has no type.", ref.name)); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Field> field, + DataTypeJsonParser::ParseType(ref.name, type_iter->value)); + ref.type = field->type(); + Result<FieldType> field_type = FieldTypeUtils::ConvertToFieldType(field); + if (field_type.ok()) { + ref.field_type = field_type.value(); + } + PAIMON_RETURN_NOT_OK(RecordReferencedType(ref.name, ref.type, types)); + return ref; +} + +bool IsDigit(char c) { + return c >= '0' && c <= '9'; +} + +/// Reads exactly `count` digits of `text` at `*pos`. +bool ReadDigits(std::string_view text, size_t* pos, size_t count, int64_t* value) { + if (text.size() - *pos < count) { + return false; + } + int64_t result = 0; + for (size_t i = 0; i < count; ++i) { + if (!IsDigit(text[*pos + i])) { + return false; + } + result = result * 10 + (text[*pos + i] - '0'); + } + *pos += count; + *value = result; + return true; +} + +bool ReadChar(std::string_view text, size_t* pos, char expected) { + if (*pos >= text.size() || text[*pos] != expected) { + return false; + } + ++*pos; + return true; +} + +/// Days from 1970-01-01 to a date of the proleptic Gregorian calendar, as Java +/// `LocalDate.toEpochDay()`. +int64_t EpochDay(int64_t year, int64_t month, int64_t day) { + const int64_t y = month <= 2 ? year - 1 : year; + const int64_t era = (y >= 0 ? y : y - 399) / 400; + const int64_t year_of_era = y - era * 400; + const int64_t day_of_year = (153 * (month > 2 ? month - 3 : month + 9) + 2) / 5 + day - 1; + const int64_t day_of_era = + year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + return era * 146097 + day_of_era - 719468; +} + +/// Reads a date as Java `DateTimeFormatter.ISO_LOCAL_DATE` parses it: a year of four digits, of +/// more only after a `+`, or negative after a `-`, then the `-MM-dd` of a real date. +bool ReadIsoDate(std::string_view text, size_t* pos, int64_t* epoch_day) { + const bool positive = *pos < text.size() && text[*pos] == '+'; + const bool negative = *pos < text.size() && text[*pos] == '-'; + if (positive || negative) { + ++*pos; + } + const size_t start = *pos; + int64_t year = 0; + while (*pos < text.size() && *pos - start < 10 && IsDigit(text[*pos])) { + year = year * 10 + (text[(*pos)++] - '0'); + } + const size_t digits = *pos - start; + if (digits < 4 || (positive && digits == 4) || (!positive && !negative && digits > 4) || + (negative && year == 0) || year > 999999999) { + return false; + } + if (negative) { + year = -year; + } + int64_t month = 0; + int64_t day = 0; + if (!ReadChar(text, pos, '-') || !ReadDigits(text, pos, 2, &month) || + !ReadChar(text, pos, '-') || !ReadDigits(text, pos, 2, &day)) { + return false; + } + static constexpr int64_t kDaysInMonth[] = {31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}; + const bool leap = (year % 4 == 0 && year % 100 != 0) || year % 400 == 0; + if (month < 1 || month > 12 || day < 1 || + day > kDaysInMonth[month - 1] + (month == 2 && leap ? 1 : 0)) { + return false; + } + *epoch_day = EpochDay(year, month, day); + return true; +} + +/// Reads a time as Java `DateTimeFormatter.ISO_LOCAL_TIME` parses it: `HH:mm`, optionally followed +/// by `:ss` and then by a fraction of one to nine digits. `seconds_required` makes only `:ss` +/// mandatory; the fraction stays optional. +bool ReadIsoTime(std::string_view text, size_t* pos, bool seconds_required, int64_t* second_of_day, + int64_t* nanos) { + int64_t hour = 0; + int64_t minute = 0; + int64_t second = 0; + *nanos = 0; + if (!ReadDigits(text, pos, 2, &hour) || !ReadChar(text, pos, ':') || + !ReadDigits(text, pos, 2, &minute) || hour > 23 || minute > 59) { + return false; + } + if (ReadChar(text, pos, ':')) { + if (!ReadDigits(text, pos, 2, &second) || second > 59) { + return false; + } + if (ReadChar(text, pos, '.')) { + size_t digits = 0; + // Accumulate at most ten digits, so an overlong fraction cannot overflow. + while (*pos < text.size() && IsDigit(text[*pos]) && digits < 10) { + *nanos = *nanos * 10 + (text[(*pos)++] - '0'); + ++digits; + } + if (digits == 0 || digits > 9) { + return false; + } + for (; digits < 9; ++digits) { + *nanos *= 10; + } + } + } else if (seconds_required) { + return false; + } + *second_of_day = hour * 3600 + minute * 60 + second; + return true; +} + +/// Parses the text Java `LocalDate.parse()` reads. +Result<int32_t> ParseDateLiteral(std::string_view text, const FieldRef& ref) { + size_t pos = 0; + int64_t epoch_day = 0; + if (!ReadIsoDate(text, &pos, &epoch_day) || pos != text.size() || + epoch_day < std::numeric_limits<int32_t>::min() || + epoch_day > std::numeric_limits<int32_t>::max()) { + return Status::Invalid( + fmt::format("Cannot parse '{}' as a date literal of column '{}'.", text, ref.name)); + } + return static_cast<int32_t>(epoch_day); +} + +/// Parses the ISO text Java `LocalDateTime.toString()` writes for a TIMESTAMP column, or the +/// text `Instant.toString()` writes with a `Z` for a TIMESTAMP_LTZ one. Expanded years are +/// supported; the instant must lie within the range of epoch milliseconds. +Result<Literal> ParseTimestampLiteral(std::string_view text, const FieldRef& ref) { + const bool instant = !checked_cast<const arrow::TimestampType&>(*ref.type).timezone().empty(); + size_t pos = 0; + int64_t epoch_day = 0; + int64_t second_of_day = 0; + int64_t nanos = 0; + if (!ReadIsoDate(text, &pos, &epoch_day) || !ReadChar(text, &pos, 'T') || + !ReadIsoTime(text, &pos, /*seconds_required=*/instant, &second_of_day, &nanos) || + (instant && !ReadChar(text, &pos, 'Z')) || pos != text.size()) { + return Status::Invalid(fmt::format( + "Cannot parse '{}' as a timestamp literal of column '{}'.", text, ref.name)); + } + const int64_t seconds = epoch_day * 86400 + second_of_day; + // Check the complete value: the fractional milliseconds can bring the first second into + // range, and only part of the last second fits. Widen before multiplying to avoid overflow. + const Decimal::int128_t millis = + static_cast<Decimal::int128_t>(seconds) * 1000 + nanos / 1000000; + if (millis > std::numeric_limits<int64_t>::max() || + millis < std::numeric_limits<int64_t>::min()) { + return Status::Invalid(fmt::format( + "Timestamp literal '{}' of column '{}' is outside the range of epoch milliseconds.", + text, ref.name)); + } + return Literal(Timestamp::FromEpochMillis(static_cast<int64_t>(millis), + static_cast<int32_t>(nanos % 1000000))); +} + +Result<Literal> ParseDecimalLiteral(const std::string& text, const FieldRef& ref) { + const auto& decimal_type = checked_cast<const arrow::DecimalType&>(*ref.type); + arrow::Decimal128 value; + int32_t precision = 0; + int32_t scale = 0; + if (!arrow::Decimal128::FromString(text, &value, &precision, &scale).ok()) { + return Status::Invalid( + fmt::format("Cannot parse '{}' as a decimal literal of column '{}'.", text, ref.name)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(arrow::Decimal128 rescaled, + value.Rescale(scale, decimal_type.scale())); + if (!rescaled.FitsInPrecision(decimal_type.precision())) { + return Status::Invalid( + fmt::format("Decimal literal '{}' does not fit column '{}' of type {}.", text, ref.name, + ref.type->ToString())); + } + return Literal(Decimal( + decimal_type.precision(), decimal_type.scale(), + static_cast<Decimal::int128_t>( + static_cast<Decimal::uint128_t>(static_cast<uint64_t>(rescaled.high_bits())) << 64 | + rescaled.low_bits()))); +} + +/// Parses one literal as Java `LeafPredicate` writes it for a column of the type of `ref`. +Result<Literal> ParseLiteral(const rapidjson::Value& value, const rapidjson::Value& number_text, + const FieldRef& ref) { + if (value.IsNull()) { + return Literal(ref.field_type); + } + switch (ref.field_type) { + case FieldType::BOOLEAN: + if (value.IsBool()) { + return Literal(value.GetBool()); + } + break; + case FieldType::TINYINT: + if (value.IsInt() && value.GetInt() >= std::numeric_limits<int8_t>::min() && + value.GetInt() <= std::numeric_limits<int8_t>::max()) { + return Literal(static_cast<int8_t>(value.GetInt())); + } + break; + case FieldType::SMALLINT: + if (value.IsInt() && value.GetInt() >= std::numeric_limits<int16_t>::min() && + value.GetInt() <= std::numeric_limits<int16_t>::max()) { + return Literal(static_cast<int16_t>(value.GetInt())); + } + break; + case FieldType::INT: + if (value.IsInt()) { + return Literal(static_cast<int32_t>(value.GetInt())); + } + break; + case FieldType::BIGINT: + if (value.IsInt64()) { + return Literal(static_cast<int64_t>(value.GetInt64())); + } + break; + case FieldType::FLOAT: + // Preserve integer-to-float rounding as Java Number.floatValue() does; converting + // through double can lose the bit that decides which adjacent float to choose. + if (value.IsInt64()) { + return Literal(static_cast<float>(value.GetInt64())); + } + if (value.IsUint64()) { + return Literal(static_cast<float>(value.GetUint64())); + } + if (value.IsNumber()) { + // Java reads integers outside the long range as BigInteger. Preserve their + // spelling: rounding through double first can select the other adjacent float. + const std::string text(number_text.GetString(), number_text.GetStringLength()); + if (text.find_first_of(".eE") == std::string::npos) { + const std::optional<float> number = StringUtils::StringToValue<float>(text); + if (number && std::isfinite(number.value())) { + return Literal(number.value()); + } + break; + } + const double number = value.GetDouble(); + const double maximum = std::numeric_limits<float>::max(); + const double overflow_threshold = + maximum + std::ldexp(1.0, std::numeric_limits<float>::max_exponent - + std::numeric_limits<float>::digits - 1); + // Java writes Float.MAX_VALUE as 3.4028235E38, slightly above its exact value. + // Clamp values rounding back to that bound before the C++ narrowing conversion. + if (number > -overflow_threshold && number < overflow_threshold) { + return Literal(static_cast<float>(std::clamp(number, -maximum, maximum))); + } + } + break; + case FieldType::DOUBLE: + if (value.IsNumber()) { + return Literal(value.GetDouble()); + } + break; + case FieldType::STRING: + if (value.IsString()) { + return Literal(FieldType::STRING, value.GetString(), value.GetStringLength()); + } + break; + case FieldType::DATE: + // Not a number: Java `PredicateBuilder.convertJavaObject()` refuses one for a DATE. + if (value.IsString()) { + PAIMON_ASSIGN_OR_RAISE( + int32_t days, + ParseDateLiteral(std::string_view(value.GetString(), value.GetStringLength()), + ref)); + return Literal(FieldType::DATE, days); + } + break; + case FieldType::TIMESTAMP: + if (value.IsString()) { + return ParseTimestampLiteral( + std::string_view(value.GetString(), value.GetStringLength()), ref); + } + break; + case FieldType::DECIMAL: + if (value.IsString()) { + return ParseDecimalLiteral(std::string(value.GetString(), value.GetStringLength()), + ref); + } + break; + default: + return Status::NotImplemented( + fmt::format("Query auth row filters on column '{}' of type {} are not supported.", + ref.name, ref.type->ToString())); + } + return Status::Invalid( + fmt::format("A literal of a query auth row filter does not fit column '{}' of type {}.", + ref.name, ref.type->ToString())); +} + +/// A predicate no row satisfies, kept on a field so it binds like any other leaf. +Result<std::shared_ptr<Predicate>> AlwaysFalse(const FieldRef& ref) { + return PredicateBuilder::And( + {PredicateBuilder::IsNull(ref.index, ref.name, ref.field_type), + PredicateBuilder::IsNotNull(ref.index, ref.name, ref.field_type)}); +} + +/// Maps a Java `LeafFunction` name to a predicate; a null literal never matches, as in Java. +Result<std::shared_ptr<Predicate>> MakeLeafPredicate(const std::string& function, + const FieldRef& ref, + const std::vector<Literal>& literals) { + auto expect_literals = [&](size_t count) -> Status { + if (literals.size() != count) { + return Status::Invalid( + fmt::format("Query auth row filter {} on column '{}' needs {} literals but has {}.", + function, ref.name, count, literals.size())); + } + return Status::OK(); + }; + auto any_null = [&]() { + return std::any_of(literals.begin(), literals.end(), + [](const Literal& literal) { return literal.IsNull(); }); + }; + const int32_t index = ref.index; + const std::string& name = ref.name; + const FieldType type = ref.field_type; + if (function == "IS_NULL") { + return PredicateBuilder::IsNull(index, name, type); + } + if (function == "IS_NOT_NULL") { + return PredicateBuilder::IsNotNull(index, name, type); + } + if (function == "IN") { + // `x IN (1, NULL)` holds exactly when `x IN (1)` does. + std::vector<Literal> non_null; + for (const Literal& literal : literals) { + if (!literal.IsNull()) { + non_null.push_back(literal); + } + } + if (non_null.empty()) { + return AlwaysFalse(ref); + } + return PredicateBuilder::In(index, name, type, non_null); + } + if (function == "NOT_IN") { + // `x NOT IN (1, NULL)` never holds. + if (any_null()) { + return AlwaysFalse(ref); + } + if (literals.empty()) { + return PredicateBuilder::IsNotNull(index, name, type); + } + return PredicateBuilder::NotIn(index, name, type, literals); + } + if (function == "BETWEEN" || function == "NOT_BETWEEN") { + PAIMON_RETURN_NOT_OK(expect_literals(2)); + if (any_null()) { + return AlwaysFalse(ref); + } + if (function == "BETWEEN") { + return PredicateBuilder::Between(index, name, type, literals[0], literals[1]); + } + return PredicateBuilder::Or( + {PredicateBuilder::LessThan(index, name, type, literals[0]), + PredicateBuilder::GreaterThan(index, name, type, literals[1])}); + } + static const std::set<std::string> kBinaryFunctions = {"EQUAL", "NOT_EQUAL", + "LESS_THAN", "LESS_OR_EQUAL", + "GREATER_THAN", "GREATER_OR_EQUAL", + "STARTS_WITH", "ENDS_WITH", + "CONTAINS", "LIKE"}; + if (kBinaryFunctions.count(function) == 0) { + return Status::NotImplemented( + fmt::format("Query auth row filter function '{}' is not supported.", function)); + } + PAIMON_RETURN_NOT_OK(expect_literals(1)); + if (literals[0].IsNull()) { + return AlwaysFalse(ref); + } + const Literal& literal = literals[0]; + if (function == "EQUAL") { + return PredicateBuilder::Equal(index, name, type, literal); + } + if (function == "NOT_EQUAL") { + return PredicateBuilder::NotEqual(index, name, type, literal); + } + if (function == "LESS_THAN") { + return PredicateBuilder::LessThan(index, name, type, literal); + } + if (function == "LESS_OR_EQUAL") { + return PredicateBuilder::LessOrEqual(index, name, type, literal); + } + if (function == "GREATER_THAN") { + return PredicateBuilder::GreaterThan(index, name, type, literal); + } + if (function == "GREATER_OR_EQUAL") { + return PredicateBuilder::GreaterOrEqual(index, name, type, literal); + } + if (function == "STARTS_WITH") { + return PredicateBuilder::StartsWith(index, name, type, literal); + } + if (function == "ENDS_WITH") { + return PredicateBuilder::EndsWith(index, name, type, literal); + } + if (function == "CONTAINS") { + return PredicateBuilder::Contains(index, name, type, literal); + } + return PredicateBuilder::Like(index, name, type, literal); +} + +/// A parsed row filter: `predicate`, or, when it reads no column, the constant `value`. +struct ParsedFilter { + std::shared_ptr<Predicate> predicate; + bool value = true; +}; + +/// Combines `children` with AND, or with OR when `is_and` is false, folding away their constants +/// as Java `PredicateBuilder.and()` and `or()` do. +Result<ParsedFilter> Combine(const std::vector<ParsedFilter>& children, bool is_and) { + std::vector<std::shared_ptr<Predicate>> predicates; + for (const ParsedFilter& child : children) { + if (child.predicate != nullptr) { + predicates.push_back(child.predicate); + } else if (child.value != is_and) { + // FALSE decides an AND, TRUE an OR. + return ParsedFilter{nullptr, child.value}; + } + } + if (predicates.empty()) { + return ParsedFilter{nullptr, is_and}; + } + if (predicates.size() == 1) { + return ParsedFilter{predicates[0], true}; + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<Predicate> combined, + is_and ? PredicateBuilder::And(predicates) : PredicateBuilder::Or(predicates)); + return ParsedFilter{std::move(combined), true}; +} + +/// Parses a Java `Predicate`: `{"kind": "LEAF", "transform": ..., "function": ..., "literals": +/// [...]}` or `{"kind": "COMPOUND", "function": "AND", "children": [...]}`. +Result<ParsedFilter> ParsePredicate(const rapidjson::Value& value, + const rapidjson::Value& number_texts, ReferencedTypes* types) { + if (!value.IsObject()) { + return Status::Invalid("A query auth row filter must be a JSON object."); + } + PAIMON_ASSIGN_OR_RAISE(std::string kind, GetString(value, "kind", "A row filter")); + PAIMON_ASSIGN_OR_RAISE(std::string function, GetString(value, "function", "A row filter")); + if (kind == "COMPOUND") { + if (function != "AND" && function != "OR") { + return Status::NotImplemented( + fmt::format("Query auth row filter function '{}' is not supported.", function)); + } + auto children_iter = value.FindMember("children"); + if (children_iter == value.MemberEnd() || !children_iter->value.IsArray() || + children_iter->value.Empty()) { + return Status::Invalid("A compound row filter needs a non-empty 'children' array."); + } + // Every child is parsed before the constants fold, so one this library cannot apply fails + // and the columns each names stay recorded for the schema checks. + std::vector<ParsedFilter> children; + for (rapidjson::SizeType i = 0; i < children_iter->value.Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE( + ParsedFilter parsed, + ParsePredicate(children_iter->value[i], number_texts["children"][i], types)); + children.push_back(std::move(parsed)); + } + return Combine(children, /*is_and=*/function == "AND"); + } + if (kind != "LEAF") { + return Status::NotImplemented( + fmt::format("Query auth row filter kind '{}' is not supported.", kind)); + } + auto transform_iter = value.FindMember("transform"); + if (transform_iter == value.MemberEnd() || !transform_iter->value.IsObject()) { + return Status::Invalid("A leaf row filter needs a 'transform' object."); + } + const rapidjson::Value& transform = transform_iter->value; + PAIMON_ASSIGN_OR_RAISE(std::string transform_name, + GetString(transform, "name", "A row filter transform")); + // Checked before the constants, which ignore the literals: Java reads them as a list. + auto literals_iter = value.FindMember("literals"); + const bool has_literals = literals_iter != value.MemberEnd() && !literals_iter->value.IsNull(); + if (has_literals && !literals_iter->value.IsArray()) { + return Status::Invalid("The 'literals' of a row filter must be a JSON array."); + } + // A missing NOT_IN list must not become an empty set accepting every non-null value. + if (!has_literals && (function == "IN" || function == "NOT_IN")) { + return Status::Invalid(fmt::format( + "Query auth row filter {} requires a non-null 'literals' array.", function)); + } + const bool constant = function == "TRUE" || function == "FALSE"; + // Java `PredicateBuilder.alwaysTrue()` and `alwaysFalse()` put their constant on a NULL + // transform; the constant ignores whatever its transform computes. + if (constant && transform_name == "NULL") { + return ParsedFilter{nullptr, function == "TRUE"}; + } + if (transform_name != "FIELD_REF") { + return Status::NotImplemented(fmt::format( + "Query auth row filters on a '{}' transform are not supported.", transform_name)); + } + auto field_ref_iter = transform.FindMember("fieldRef"); + if (field_ref_iter == transform.MemberEnd()) { + return Status::Invalid("A FIELD_REF transform needs a 'fieldRef' member."); + } + PAIMON_ASSIGN_OR_RAISE(FieldRef ref, ParseFieldRef(field_ref_iter->value, types)); + if (constant) { + return ParsedFilter{nullptr, function == "TRUE"}; + } + if (ref.field_type == FieldType::UNKNOWN) { + return Status::NotImplemented( + fmt::format("Query auth row filter {} on column '{}' of type {} is not supported.", + function, ref.name, ref.type->ToString())); + } + std::vector<Literal> literals; + if (has_literals) { + for (rapidjson::SizeType i = 0; i < literals_iter->value.Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(Literal parsed, ParseLiteral(literals_iter->value[i], + number_texts["literals"][i], ref)); + literals.push_back(std::move(parsed)); + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> predicate, + MakeLeafPredicate(function, ref, literals)); + return ParsedFilter{std::move(predicate), true}; +} + +Result<rapidjson::Document> ParseJson(const std::string& json, const std::string& what) { + // An unescaped NUL terminates RapidJSON's input before it validates the remaining bytes. + if (json.find('\0') != std::string::npos) { + return Status::Invalid(fmt::format("{} contains an unescaped NUL byte.", what)); + } + rapidjson::Document document; + // Approximate double parsing can shift a comparison boundary by a few ULPs and authorize + // different rows than Java. Preserve the correctly rounded value of each numeric literal. + document.Parse<rapidjson::kParseValidateEncodingFlag | rapidjson::kParseFullPrecisionFlag>( + json.c_str(), json.size()); + if (document.HasParseError()) { + return Status::Invalid( + fmt::format("{} is not valid JSON at offset {}.", what, document.GetErrorOffset())); + } + Status unique = RapidJsonUtil::CheckNoDuplicateMembers(document); + if (!unique.ok()) { + return Status::Invalid(fmt::format("{} is ambiguous: {}.", what, unique.message())); + } + return document; +} + +Result<std::shared_ptr<arrow::Array>> FindColumn(const arrow::StructArray& batch, + const std::string& name) { + std::shared_ptr<arrow::Array> column = batch.GetFieldByName(name); + if (column == nullptr) { + return Status::Invalid( + fmt::format("Column '{}', which a column mask reads, is not in the read batch.", name)); + } + return column; +} + +Status CheckStringTarget(const std::string& transform, const DataField& target) { + if (target.Type()->id() != arrow::Type::STRING) { + return Status::Invalid( + fmt::format("Column masking {} produces a string, but column '{}' is {}.", transform, + target.Name(), target.Type()->ToString())); + } + return Status::OK(); +} + +class NullMask : public MaskTransform { + public: + const std::string& Name() const override { + static const std::string kName = "NULL"; + return kName; + } + std::vector<std::string> InputFields() const override { + return {}; + } + Status ValidateTarget(const DataField& target) const override { + return Status::OK(); + } + Result<std::shared_ptr<arrow::Array>> Evaluate( + const arrow::StructArray& batch, const std::shared_ptr<arrow::DataType>& target_type, + arrow::MemoryPool* pool) const override { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> nulls, + arrow::MakeArrayOfNull(target_type, batch.length(), pool)); + return nulls; + } +}; + +class FieldRefMask : public MaskTransform { + public: + explicit FieldRefMask(FieldRef source) : source_(std::move(source)) {} + + const std::string& Name() const override { + static const std::string kName = "FIELD_REF"; + return kName; + } + std::vector<std::string> InputFields() const override { + return {source_.name}; + } + Status ValidateTarget(const DataField& target) const override { + if (!target.Type()->Equals(*source_.type)) { + return Status::Invalid(fmt::format( + "Column masking on '{}' copies column '{}' of type {} into a column of type {}.", + target.Name(), source_.name, source_.type->ToString(), target.Type()->ToString())); + } + return Status::OK(); + } + Result<std::shared_ptr<arrow::Array>> Evaluate( + const arrow::StructArray& batch, const std::shared_ptr<arrow::DataType>& target_type, + arrow::MemoryPool* pool) const override { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> column, + FindColumn(batch, source_.name)); + if (!column->type()->Equals(*target_type)) { + return Status::Invalid(fmt::format( + "Column masking copies column '{}' of type {} into a column of type {}.", + source_.name, column->type()->ToString(), target_type->ToString())); + } + return column; + } + + private: + FieldRef source_; +}; + +struct StringOperand { + enum class Kind { kNull, kLiteral, kField }; + Kind kind = Kind::kNull; + std::string literal; + std::string field; +}; + +Result<StringOperand> ParseStringOperand(const rapidjson::Value& value, ReferencedTypes* types) { + StringOperand operand; + if (value.IsNull()) { + return operand; + } + if (value.IsString()) { + operand.kind = StringOperand::Kind::kLiteral; + operand.literal.assign(value.GetString(), value.GetStringLength()); + return operand; + } + if (value.IsObject()) { + PAIMON_ASSIGN_OR_RAISE(FieldRef ref, ParseFieldRef(value, types)); + if (ref.type->id() != arrow::Type::STRING) { + return Status::Invalid(fmt::format( + "A string transform of a column mask reads column '{}' of type {}, not a string.", + ref.name, ref.type->ToString())); + } + operand.kind = StringOperand::Kind::kField; + operand.field = ref.name; + return operand; + } + return Status::Invalid("A string transform input must be a string, a field reference or null."); +} + +class StringOperandReader { + public: + static Result<StringOperandReader> Create(const StringOperand& operand, + const arrow::StructArray& batch) { + StringOperandReader reader(operand); + if (operand.kind == StringOperand::Kind::kField) { + PAIMON_ASSIGN_OR_RAISE(reader.column_, FindColumn(batch, operand.field)); + if (reader.column_->type_id() != arrow::Type::STRING) { + return Status::Invalid(fmt::format("Column '{}' a string mask reads is {}.", + operand.field, + reader.column_->type()->ToString())); + } + } + return reader; + } + + std::optional<std::string_view> Get(int64_t row) const { + switch (operand_.kind) { + case StringOperand::Kind::kNull: + return std::nullopt; + case StringOperand::Kind::kLiteral: + return std::string_view(operand_.literal); + case StringOperand::Kind::kField: { + const auto& strings = checked_cast<const arrow::StringArray&>(*column_); + if (strings.IsNull(row)) { + return std::nullopt; + } + return strings.GetView(row); + } + } + return std::nullopt; + } + + private: + explicit StringOperandReader(const StringOperand& operand) : operand_(operand) {} + + const StringOperand& operand_; + std::shared_ptr<arrow::Array> column_; +}; + +/// `CONCAT` is null when any input is; `CONCAT_WS` is null when its separator is and skips the +/// other null inputs, as Java `BinaryString` does. +class ConcatMask : public MaskTransform { + public: + ConcatMask(std::vector<StringOperand> inputs, bool with_separator) + : inputs_(std::move(inputs)), with_separator_(with_separator) {} + + const std::string& Name() const override { + static const std::string kConcat = "CONCAT"; + static const std::string kConcatWs = "CONCAT_WS"; + return with_separator_ ? kConcatWs : kConcat; + } + std::vector<std::string> InputFields() const override { + std::vector<std::string> fields; + for (const StringOperand& input : inputs_) { + if (input.kind == StringOperand::Kind::kField) { + fields.push_back(input.field); + } + } + return fields; + } + Status ValidateTarget(const DataField& target) const override { + return CheckStringTarget(Name(), target); + } + Result<std::shared_ptr<arrow::Array>> Evaluate( + const arrow::StructArray& batch, const std::shared_ptr<arrow::DataType>& target_type, + arrow::MemoryPool* pool) const override { + std::vector<StringOperandReader> readers; + readers.reserve(inputs_.size()); + for (const StringOperand& input : inputs_) { + PAIMON_ASSIGN_OR_RAISE(StringOperandReader reader, + StringOperandReader::Create(input, batch)); + readers.push_back(std::move(reader)); + } + arrow::StringBuilder builder(pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(batch.length())); + std::string value; + for (int64_t row = 0; row < batch.length(); ++row) { + value.clear(); + bool is_null = false; + if (with_separator_) { + std::optional<std::string_view> separator = readers[0].Get(row); + is_null = !separator.has_value(); + bool first = true; + for (size_t i = 1; !is_null && i < readers.size(); ++i) { + std::optional<std::string_view> input = readers[i].Get(row); + if (!input) { + continue; + } + if (!first) { + value.append(separator->data(), separator->size()); + } + value.append(input->data(), input->size()); + first = false; + } + } else { + for (const StringOperandReader& reader : readers) { + std::optional<std::string_view> input = reader.Get(row); + if (!input) { + is_null = true; + break; + } + value.append(input->data(), input->size()); + } + } + if (is_null) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(value)); + } + } + std::shared_ptr<arrow::Array> result; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&result)); + return result; + } + + private: + std::vector<StringOperand> inputs_; + bool with_separator_; +}; + +struct PositionOperand { + enum class Kind { kNull, kLiteral, kField }; + Kind kind = Kind::kNull; + int32_t literal = 0; + std::string field; +}; + +Result<PositionOperand> ParsePositionOperand(const rapidjson::Value& value, + ReferencedTypes* types) { + PositionOperand operand; + if (value.IsNull()) { + return operand; + } + if (value.IsInt()) { + operand.kind = PositionOperand::Kind::kLiteral; + operand.literal = value.GetInt(); + return operand; + } + if (value.IsObject()) { + PAIMON_ASSIGN_OR_RAISE(FieldRef ref, ParseFieldRef(value, types)); + if (!FieldTypeUtils::IsIntegerNumeric(ref.field_type)) { + return Status::Invalid(fmt::format( + "SUBSTRING of a column mask takes a position from column '{}' of type {}, not an " + "integer.", + ref.name, ref.type->ToString())); + } + operand.kind = PositionOperand::Kind::kField; + operand.field = ref.name; + return operand; + } + return Status::Invalid("A SUBSTRING position must be an integer, a field reference or null."); +} + +/// Java `BinaryString.substringSQL`: `pos` counts characters from 1, or from the end when negative. +std::string_view SubstringSql(std::string_view value, int32_t pos, int32_t length) { + int64_t num_chars = 0; + for (size_t i = 0; i < value.size(); i += BinaryString::NumBytesForFirstByte(value[i])) { + ++num_chars; + } + int64_t start = 0; + if (pos > 0) { + start = static_cast<int64_t>(pos) - 1; + } else if (pos < 0) { + start = num_chars + pos; + } + const int64_t end = std::clamp<int64_t>(start + length, std::numeric_limits<int32_t>::min(), + std::numeric_limits<int32_t>::max()); + // Java `BinaryString.substring`, comparing a character index with the size in bytes as it does. + if (end <= start || start >= static_cast<int64_t>(value.size())) { + return {}; + } + size_t i = 0; + int64_t c = 0; + while (i < value.size() && c < start) { + i += BinaryString::NumBytesForFirstByte(value[i]); + ++c; + } + const size_t begin = std::min(i, value.size()); + while (i < value.size() && c < end) { + i += BinaryString::NumBytesForFirstByte(value[i]); + ++c; + } + i = std::min(i, value.size()); + return i > begin ? value.substr(begin, i - begin) : std::string_view(); +} + +/// `SUBSTRING(source, pos[, length])`, null when any input is. +class SubstringMask : public MaskTransform { + public: + /// @param positions `pos`, then the optional `length`. + SubstringMask(StringOperand source, std::vector<PositionOperand> positions) + : source_(std::move(source)), positions_(std::move(positions)) {} + + const std::string& Name() const override { + static const std::string kName = "SUBSTRING"; + return kName; + } + std::vector<std::string> InputFields() const override { + std::vector<std::string> fields; + if (source_.kind == StringOperand::Kind::kField) { + fields.push_back(source_.field); + } + for (const PositionOperand& position : positions_) { + if (position.kind == PositionOperand::Kind::kField) { + fields.push_back(position.field); + } + } + return fields; + } + Status ValidateTarget(const DataField& target) const override { + return CheckStringTarget(Name(), target); + } + Result<std::shared_ptr<arrow::Array>> Evaluate( + const arrow::StructArray& batch, const std::shared_ptr<arrow::DataType>& target_type, + arrow::MemoryPool* pool) const override { + PAIMON_ASSIGN_OR_RAISE(StringOperandReader source, + StringOperandReader::Create(source_, batch)); + std::vector<std::shared_ptr<arrow::Array>> position_columns(positions_.size()); + for (size_t i = 0; i < positions_.size(); ++i) { + if (positions_[i].kind == PositionOperand::Kind::kField) { + PAIMON_ASSIGN_OR_RAISE(position_columns[i], FindColumn(batch, positions_[i].field)); + } + } + arrow::StringBuilder builder(pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(batch.length())); + for (int64_t row = 0; row < batch.length(); ++row) { + std::optional<std::string_view> value = source.Get(row); + bool null_position = false; + for (size_t i = 0; i < positions_.size(); ++i) { + if (positions_[i].kind == PositionOperand::Kind::kNull || + (positions_[i].kind == PositionOperand::Kind::kField && + position_columns[i]->IsNull(row))) { + null_position = true; + break; + } + } + // Java checks every input for NULL before converting either position to an int. + if (!value || null_position) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + continue; + } + std::optional<int32_t> pos; + std::optional<int32_t> length = std::numeric_limits<int32_t>::max(); + PAIMON_ASSIGN_OR_RAISE(pos, ReadPosition(positions_[0], position_columns[0], row)); + if (positions_.size() == 2) { + PAIMON_ASSIGN_OR_RAISE(length, + ReadPosition(positions_[1], position_columns[1], row)); + } + PAIMON_RETURN_NOT_OK_FROM_ARROW( + builder.Append(SubstringSql(value.value(), pos.value(), length.value()))); + } + std::shared_ptr<arrow::Array> result; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&result)); + return result; + } + + private: + static Result<std::optional<int32_t>> ReadPosition(const PositionOperand& position, + const std::shared_ptr<arrow::Array>& column, + int64_t row) { + if (position.kind == PositionOperand::Kind::kNull) { + return std::optional<int32_t>(); + } + if (position.kind == PositionOperand::Kind::kLiteral) { + return std::optional<int32_t>(position.literal); + } + if (column->IsNull(row)) { + return std::optional<int32_t>(); + } + int64_t value = 0; + switch (column->type_id()) { + case arrow::Type::INT8: + value = checked_cast<const arrow::Int8Array&>(*column).Value(row); + break; + case arrow::Type::INT16: + value = checked_cast<const arrow::Int16Array&>(*column).Value(row); + break; + case arrow::Type::INT32: + value = checked_cast<const arrow::Int32Array&>(*column).Value(row); + break; + case arrow::Type::INT64: + value = checked_cast<const arrow::Int64Array&>(*column).Value(row); + break; + default: + return Status::Invalid(fmt::format("Column '{}' a SUBSTRING mask reads is {}.", + position.field, column->type()->ToString())); + } + if (value < std::numeric_limits<int32_t>::min() || + value > std::numeric_limits<int32_t>::max()) { + return Status::Invalid(fmt::format( + "SUBSTRING position of column '{}' does not fit an int.", position.field)); + } + return std::optional<int32_t>(static_cast<int32_t>(value)); + } + + StringOperand source_; + std::vector<PositionOperand> positions_; +}; + +/// Parses a Java `Transform` used as a column mask: `{"name": "CONCAT", "inputs": [...]}`. +Result<std::shared_ptr<MaskTransform>> ParseMaskTransform(const rapidjson::Value& value, + ReferencedTypes* types) { + if (value.IsNull()) { + return Status::Invalid("Column mask transform cannot be JSON null."); + } + if (!value.IsObject()) { + return Status::Invalid("A column mask transform must be a JSON object."); + } + PAIMON_ASSIGN_OR_RAISE(std::string name, GetString(value, "name", "A column mask transform")); + if (name == "NULL") { + return std::make_shared<NullMask>(); + } + if (name == "FIELD_REF") { + auto field_ref_iter = value.FindMember("fieldRef"); + if (field_ref_iter == value.MemberEnd()) { + return Status::Invalid("A FIELD_REF transform needs a 'fieldRef' member."); + } + PAIMON_ASSIGN_OR_RAISE(FieldRef ref, ParseFieldRef(field_ref_iter->value, types)); + return std::make_shared<FieldRefMask>(std::move(ref)); + } + if (name != "CONCAT" && name != "CONCAT_WS" && name != "SUBSTRING") { + return Status::NotImplemented( + fmt::format("Column mask transform '{}' is not supported.", name)); + } + auto inputs_iter = value.FindMember("inputs"); + if (inputs_iter == value.MemberEnd() || !inputs_iter->value.IsArray()) { + return Status::Invalid(fmt::format("A {} transform needs an 'inputs' array.", name)); + } + const auto& inputs = inputs_iter->value; + if (name == "SUBSTRING") { + if (inputs.Size() != 2 && inputs.Size() != 3) { + return Status::Invalid("A SUBSTRING transform takes 2 or 3 inputs."); + } + PAIMON_ASSIGN_OR_RAISE(StringOperand source, ParseStringOperand(*inputs.Begin(), types)); + std::vector<PositionOperand> positions; + for (rapidjson::SizeType i = 1; i < inputs.Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(PositionOperand position, + ParsePositionOperand(inputs[i], types)); + positions.push_back(std::move(position)); + } + return std::make_shared<SubstringMask>(std::move(source), std::move(positions)); + } + if (name == "CONCAT_WS" && inputs.Size() < 2) { + return Status::Invalid("A CONCAT_WS transform needs a separator and at least one input."); + } + std::vector<StringOperand> operands; + for (const rapidjson::Value& input : inputs.GetArray()) { + PAIMON_ASSIGN_OR_RAISE(StringOperand operand, ParseStringOperand(input, types)); + operands.push_back(std::move(operand)); + } + return std::make_shared<ConcatMask>(std::move(operands), name == "CONCAT_WS"); +} + +/// Java `SpecialFields.SYSTEM_FIELD_NAMES`, none of which a query-auth read projects. +bool IsSystemColumn(const std::string& name) { + return name == SpecialFields::SequenceNumber().Name() || + name == SpecialFields::ValueKind().Name() || name == "_LEVEL" || + name == SpecialFields::RowKind().Name() || name == SpecialFields::RowId().Name(); +} + +const DataField* FindField(const std::vector<DataField>& fields, const std::string& name) { + for (const DataField& field : fields) { + if (field.Name() == name) { + return &field; + } + } + return nullptr; +} + +const DataField* FindFieldById(const std::vector<DataField>& fields, int32_t id) { + for (const DataField& field : fields) { + if (field.Id() == id) { + return &field; + } + } + return nullptr; +} + +std::vector<std::string> FieldNames(const std::vector<DataField>& fields) { + std::vector<std::string> names; + names.reserve(fields.size()); + for (const DataField& field : fields) { + names.push_back(field.Name()); + } + return names; +} + +} // namespace + +Result<std::shared_ptr<TableQueryAuthResult>> TableQueryAuthResult::Create( + const std::vector<std::string>& filter, + const std::map<std::string, std::string>& column_masking) { + std::shared_ptr<TableQueryAuthResult> result(new TableQueryAuthResult(filter, column_masking)); + std::vector<ParsedFilter> filters; + for (const std::string& json : filter) { + if (json.empty()) { + return Status::Invalid("Row filter cannot be empty."); + } + PAIMON_ASSIGN_OR_RAISE(rapidjson::Document document, ParseJson(json, "Row filter")); + if (document.IsNull()) { + return Status::Invalid("Row filter cannot be JSON null."); + } + // A second view retains numeric token spellings without treating quoted literals as + // numbers. FLOAT conversion needs these for integers outside the 64-bit range. + rapidjson::Document number_texts; + number_texts.Parse<rapidjson::kParseNumbersAsStringsFlag>(json.c_str(), json.size()); + if (number_texts.HasParseError()) { + return Status::Invalid("Row filter numeric literals could not be parsed."); + } + PAIMON_ASSIGN_OR_RAISE(ParsedFilter parsed, + ParsePredicate(document, number_texts, &result->referenced_types_)); + filters.push_back(std::move(parsed)); + } + PAIMON_ASSIGN_OR_RAISE(ParsedFilter row_filter, Combine(filters, /*is_and=*/true)); + result->rejects_all_rows_ = row_filter.predicate == nullptr && !row_filter.value; + result->row_filter_ = std::move(row_filter.predicate); + if (result->row_filter_ != nullptr) { + PAIMON_RETURN_NOT_OK( + PredicateUtils::GetAllNames(result->row_filter_, &result->row_filter_fields_)); + } + for (const auto& [column, json] : column_masking) { + if (column.empty()) { + return Status::Invalid("Column mask target cannot be empty."); + } + if (json.empty()) { + return Status::Invalid("Column mask transform cannot be empty."); + } + PAIMON_ASSIGN_OR_RAISE(rapidjson::Document document, + ParseJson(json, "Column mask transform")); + if (IsSystemColumn(column)) { + // As in Java for a system column the query does not project, the mask never reaches + // the output: it is parsed, but neither validated nor applied. + ReferencedTypes inert_types; + PAIMON_RETURN_NOT_OK(ParseMaskTransform(document, &inert_types).status()); + continue; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<MaskTransform> transform, + ParseMaskTransform(document, &result->referenced_types_)); + result->masks_.emplace(column, std::move(transform)); + } + return result; +} + +std::set<std::string> TableQueryAuthResult::MaskedFields() const { + std::set<std::string> fields; + for (const auto& [column, transform] : masks_) { + fields.insert(column); + } + return fields; +} + +bool TableQueryAuthResult::SameFilter(const TableQueryAuthResult& other) const { + if (rejects_all_rows_ != other.rejects_all_rows_) { + return false; + } + // As Java, compares the parsed conjuncts, so neither the JSON text nor their order matters. + std::vector<std::shared_ptr<Predicate>> conjuncts = PredicateUtils::SplitAnd(row_filter_); + std::vector<std::shared_ptr<Predicate>> other_conjuncts = + PredicateUtils::SplitAnd(other.row_filter_); + auto contains_all = [](const std::vector<std::shared_ptr<Predicate>>& left, + const std::vector<std::shared_ptr<Predicate>>& right) { + return std::all_of(right.begin(), right.end(), [&](const std::shared_ptr<Predicate>& r) { + return std::any_of(left.begin(), left.end(), + [&](const std::shared_ptr<Predicate>& l) { return *l == *r; }); + }); + }; + return contains_all(conjuncts, other_conjuncts) && contains_all(other_conjuncts, conjuncts); +} + +Status TableQueryAuthResult::ValidateAgainstSchema( + const std::vector<DataField>& latest_fields) const { + for (const auto& [target, json] : column_masking_) { + if (IsSystemColumn(target) && FindField(latest_fields, target) != nullptr) { + return Status::Invalid(fmt::format( + "Column mask target '{}' names a table column reserved as a system column; " + "refusing to ignore its mask.", + target)); + } + } + auto check_exists = [&](const std::string& rule, const std::string& field) -> Status { + if (FindField(latest_fields, field) == nullptr) { + return Status::Invalid(fmt::format( + "{} references column '{}' which does not exist in table schema [{}]. The rule " + "may be stale after a column rename or drop; refusing to read.", + rule, field, fmt::join(FieldNames(latest_fields), ", "))); + } + return Status::OK(); + }; + for (const std::string& field : row_filter_fields_) { + PAIMON_RETURN_NOT_OK(check_exists("Row filter", field)); + } + for (const auto& [target, transform] : masks_) { + PAIMON_RETURN_NOT_OK(check_exists("Column masking", target)); + PAIMON_RETURN_NOT_OK(transform->ValidateTarget(*FindField(latest_fields, target))); + for (const std::string& input : transform->InputFields()) { + PAIMON_RETURN_NOT_OK(check_exists("Column masking", input)); + if (input != target && masks_.count(input) > 0) { + return Status::Invalid(fmt::format( + "Column masking on '{}' reads column '{}', which is masked too. The mask " + "would be computed from the raw value of '{}' and expose it through '{}'.", + target, input, input, target)); + } + } + } + // Also every column of a branch a constant folded away, as Java validates the parsed rules. + for (const auto& [name, type] : referenced_types_) { + PAIMON_RETURN_NOT_OK(check_exists("A query auth rule", name)); + const DataField* field = FindField(latest_fields, name); + if (!field->Type()->Equals(*type)) { + return Status::Invalid(fmt::format( + "Query auth rules read column '{}' as {}, but the table holds it as {}; refusing " + "to read.", + name, type->ToString(), field->Type()->ToString())); + } + } + return Status::OK(); +} + +std::set<std::string> TableQueryAuthResult::ReferencedFields() const { + std::set<std::string> referenced = row_filter_fields_; + for (const auto& [name, type] : referenced_types_) { + referenced.insert(name); + } + for (const auto& [target, transform] : masks_) { + referenced.insert(target); + for (const std::string& input : transform->InputFields()) { + referenced.insert(input); + } + } + return referenced; +} + +std::map<std::string, int32_t> TableQueryAuthResult::BindFieldIds( + const std::vector<DataField>& fields) const { + std::map<std::string, int32_t> field_ids; + for (const std::string& name : ReferencedFields()) { + const DataField* field = FindField(fields, name); + if (field != nullptr) { + field_ids.emplace(name, field->Id()); + } + } + return field_ids; +} + +Status TableQueryAuthResult::ValidateFieldIds(const std::map<std::string, int32_t>& field_ids, + const std::vector<DataField>& fields) const { + // A column left unbound would be matched by name alone. + for (const std::string& name : ReferencedFields()) { + if (field_ids.count(name) == 0) { + return Status::Invalid(fmt::format( + "Query auth rules reference column '{}', which the schema the scan planned with " + "did not hold. Create the scan again.", + name)); + } + } + for (const auto& [name, id] : field_ids) { + const DataField* field = FindField(fields, name); + if (field != nullptr) { + if (field->Id() != id) { + return Status::Invalid(fmt::format( + "Query auth rules reference column '{}' which the schema being read exposes " + "as a different column of the same name (field {}, while the rules were bound " + "to field {}); refusing to read to avoid applying the rule to unrelated data.", + name, field->Id(), id)); + } + continue; + } + const DataField* renamed = FindFieldById(fields, id); + if (renamed != nullptr) { + return Status::Invalid(fmt::format( + "Query auth rules reference column '{}' which the schema being read exposes as " + "'{}' (renamed since), field {}; refusing to read to avoid applying the rule by a " + "stale name.", + name, renamed->Name(), id)); + } + } + return Status::OK(); +} + +std::set<std::string> TableQueryAuthResult::RequiredFields( + const std::vector<std::string>& visible_fields) const { + std::set<std::string> required = row_filter_fields_; + std::set<std::string> readable(visible_fields.begin(), visible_fields.end()); + std::deque<std::string> newly_readable(visible_fields.begin(), visible_fields.end()); + for (const std::string& field : row_filter_fields_) { + if (readable.insert(field).second) { + newly_readable.push_back(field); + } + } + while (!newly_readable.empty()) { + std::string field = std::move(newly_readable.front()); + newly_readable.pop_front(); + auto iter = masks_.find(field); + if (iter == masks_.end()) { + continue; + } + for (const std::string& input : iter->second->InputFields()) { + required.insert(input); + if (readable.insert(input).second) { + newly_readable.push_back(input); + } + } + } + return required; +} + +} // namespace paimon
diff --git a/src/paimon/core/catalog/table_query_auth_result.h b/src/paimon/core/catalog/table_query_auth_result.h new file mode 100644 index 0000000..4bfe02c --- /dev/null +++ b/src/paimon/core/catalog/table_query_auth_result.h
@@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <cstdint> +#include <map> +#include <memory> +#include <set> +#include <string> +#include <vector> + +#include "paimon/common/types/data_field.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace arrow { +class Array; +class DataType; +class MemoryPool; +class StructArray; +} // namespace arrow + +namespace paimon { +class Predicate; + +/// Computes the values replacing a masked column, mirroring the Java `Transform` of +/// `columnMasking`. Only `NULL`, `FIELD_REF`, `CONCAT`, `CONCAT_WS` and `SUBSTRING` are supported. +class MaskTransform { + public: + virtual ~MaskTransform() = default; + + virtual const std::string& Name() const = 0; + + virtual std::vector<std::string> InputFields() const = 0; + + virtual Status ValidateTarget(const DataField& target) const = 0; + + /// Computes the masked values for every row of `batch`, which holds every input column. The + /// inputs and `target_type` are plain, not dictionary-encoded, and a string is not a large one. + virtual Result<std::shared_ptr<arrow::Array>> Evaluate( + const arrow::StructArray& batch, const std::shared_ptr<arrow::DataType>& target_type, + arrow::MemoryPool* pool) const = 0; +}; + +/// The row filter and column masks of one query of a table with `query-auth.enabled=true`, +/// mirroring Java `TableQueryAuthResult`. The rules are parsed on creation, so a rule this library +/// cannot apply fails the query instead of being dropped. +class TableQueryAuthResult { + public: + /// Parses the JSON rules of an `AuthTableQueryResponse`. + static Result<std::shared_ptr<TableQueryAuthResult>> Create( + const std::vector<std::string>& filter, + const std::map<std::string, std::string>& column_masking); + + const std::vector<std::string>& Filter() const { + return filter_; + } + const std::map<std::string, std::string>& ColumnMasking() const { + return column_masking_; + } + + /// The conjunction of the filter, or null when it reads no column. Bind it to a schema by + /// name before evaluating it. + const std::shared_ptr<Predicate>& RowFilter() const { + return row_filter_; + } + /// Whether the filter folds to FALSE, as Java `PredicateBuilder.alwaysFalse()`: no row is + /// read, and `RowFilter()` is null. + bool RejectsAllRows() const { + return rejects_all_rows_; + } + /// The columns `RowFilter()` reads; see `ReferencedFields()` for all the filter names. + const std::set<std::string>& RowFilterFields() const { + return row_filter_fields_; + } + /// The masks to apply, by the column they replace. A mask of a system column such as + /// `_ROW_ID`, which a query-auth read never projects, is left out, as Java leaves it inert. + /// `ValidateAgainstSchema()` refuses a table column whose mask would be skipped this way. + const std::map<std::string, std::shared_ptr<MaskTransform>>& Masks() const { + return masks_; + } + std::set<std::string> MaskedFields() const; + + bool HasRules() const { + return row_filter_ != nullptr || rejects_all_rows_ || !masks_.empty(); + } + + /// Whether `other` holds the same row filter conjuncts, however they were serialized. + bool SameFilter(const TableQueryAuthResult& other) const; + + /// Java `validateAgainstSchema`: every column a rule reads exists with the type the rule + /// expects, a mask fits its column, and no mask reads another masked column. + Status ValidateAgainstSchema(const std::vector<DataField>& latest_fields) const; + + /// The columns the rules reference: the row filter operands, those of the branches its + /// constants folded away included, and the targets and inputs of the masks. + std::set<std::string> ReferencedFields() const; + + /// The field id of each column the rules reference, as `fields` hold them; a column `fields` + /// do not hold is left out. + std::map<std::string, int32_t> BindFieldIds(const std::vector<DataField>& fields) const; + + /// Java `validateReadableWithoutRename`: fails when `field_ids` leave a column the rules + /// reference unbound, or `fields` expose a bound column under another name, or another column + /// under its name. Only names and field ids are compared; `ValidateAgainstSchema()` checks the + /// types the rules declare. + Status ValidateFieldIds(const std::map<std::string, int32_t>& field_ids, + const std::vector<DataField>& fields) const; + + /// The columns the rules read for a read outputting `visible_fields`: the row filter operands + /// and, transitively, the inputs of every mask whose target is visible. + std::set<std::string> RequiredFields(const std::vector<std::string>& visible_fields) const; + + private: + TableQueryAuthResult(const std::vector<std::string>& filter, + const std::map<std::string, std::string>& column_masking) + : filter_(filter), column_masking_(column_masking) {} + + std::vector<std::string> filter_; + std::map<std::string, std::string> column_masking_; + std::shared_ptr<Predicate> row_filter_; + bool rejects_all_rows_ = false; + std::set<std::string> row_filter_fields_; + std::map<std::string, std::shared_ptr<MaskTransform>> masks_; + /// The type every rule was written for, per column it reads, even in a branch a constant + /// folded away. + std::map<std::string, std::shared_ptr<arrow::DataType>> referenced_types_; +}; + +} // namespace paimon
diff --git a/src/paimon/core/catalog/table_query_auth_result_test.cpp b/src/paimon/core/catalog/table_query_auth_result_test.cpp new file mode 100644 index 0000000..2112ad9 --- /dev/null +++ b/src/paimon/core/catalog/table_query_auth_result_test.cpp
@@ -0,0 +1,949 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/catalog/table_query_auth_result.h" + +#include <cmath> +#include <limits> +#include <map> +#include <memory> +#include <set> +#include <string> +#include <tuple> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "gtest/gtest.h" +#include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/predicate/leaf_predicate.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/result.h" +#include "paimon/status.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { + +std::string FieldRefJson(int32_t index, const std::string& name, const std::string& type) { + return fmt::format(R"({{"index": {}, "name": "{}", "type": "{}"}})", index, name, type); +} + +std::string LeafJson(const std::string& field_ref, const std::string& function, + const std::string& literals) { + return fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "FIELD_REF", "fieldRef": {}}},)" + R"( "function": "{}", "literals": {}}})", + field_ref, function, literals); +} + +std::string ConstantJson(const std::string& function, const std::string& literals = "[]") { + return fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "NULL"}}, "function": "{}",)" + R"( "literals": {}}})", + function, literals); +} + +std::string CompoundJson(const std::string& function, const std::string& left, + const std::string& right) { + return fmt::format(R"({{"kind": "COMPOUND", "function": "{}", "children": [{}, {}]}})", + function, left, right); +} + +Result<Literal> FilterLiteral(const std::string& field_ref, const std::string& function, + const std::string& literals) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({LeafJson(field_ref, function, literals)}, {})); + auto leaf = std::dynamic_pointer_cast<LeafPredicate>(result->RowFilter()); + if (leaf == nullptr) { + return Status::Invalid("the row filter is not a leaf"); + } + return leaf->Literals().at(0); +} + +std::string IdRef() { + return FieldRefJson(0, "id", "INT NOT NULL"); +} +std::string NameRef() { + return FieldRefJson(1, "name", "STRING"); +} +std::string PhoneRef() { + return FieldRefJson(2, "phone", "STRING"); +} + +std::vector<DataField> TableFields() { + return {DataField(0, arrow::field("id", arrow::int32(), /*nullable=*/false)), + DataField(1, arrow::field("name", arrow::utf8())), + DataField(2, arrow::field("phone", arrow::utf8())), + DataField(3, arrow::field("salary", arrow::int64()))}; +} + +std::shared_ptr<arrow::StructArray> MakeBatch(const std::string& rows) { + auto type = + arrow::struct_({arrow::field("id", arrow::int32(), /*nullable=*/false), + arrow::field("name", arrow::utf8()), arrow::field("phone", arrow::utf8()), + arrow::field("salary", arrow::int64())}); + auto array = arrow::ipc::internal::json::ArrayFromJSON(type, rows).ValueOrDie(); + return checked_pointer_cast<arrow::StructArray>(array); +} + +Result<std::vector<char>> Evaluate(const std::shared_ptr<Predicate>& predicate, + const std::shared_ptr<arrow::StructArray>& batch) { + std::map<std::string, int32_t> field_index; + for (int32_t i = 0; i < batch->num_fields(); ++i) { + field_index[batch->struct_type()->field(i)->name()] = i; + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> bound, + PredicateUtils::CreatePickedFieldFilter(predicate, field_index)); + auto filter = std::dynamic_pointer_cast<PredicateFilter>(bound); + if (filter == nullptr) { + return Status::Invalid("the bound predicate is not a filter"); + } + return filter->Test(*batch, arrow::default_memory_pool()); +} + +Result<std::shared_ptr<arrow::Array>> Mask(const TableQueryAuthResult& result, + const std::string& column, + const std::shared_ptr<arrow::StructArray>& batch, + const std::shared_ptr<arrow::DataType>& type) { + auto mask = result.Masks().find(column); + if (mask == result.Masks().end()) { + return Status::Invalid(fmt::format("no mask on column '{}'", column)); + } + return mask->second->Evaluate(*batch, type, arrow::default_memory_pool()); +} + +std::shared_ptr<arrow::Array> Strings(const std::string& json) { + return arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), json).ValueOrDie(); +} + +} // namespace + +TEST(TableQueryAuthResultTest, ParsesTheRowFilter) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {LeafJson(IdRef(), "GREATER_THAN", "[1]"), + CompoundJson("OR", LeafJson(NameRef(), "STARTS_WITH", R"(["b"])"), + LeafJson(NameRef(), "IS_NULL", "[]"))}, + {})); + ASSERT_TRUE(result->HasRules()); + ASSERT_NE(nullptr, result->RowFilter()); + ASSERT_EQ(std::set<std::string>({"id", "name"}), result->RowFilterFields()); + ASSERT_TRUE(result->Masks().empty()); + + auto batch = MakeBatch( + R"([[1, "bob", null, 1], [2, "bob", null, 2], [3, "carol", null, 3], [4, null, null, 4]])"); + ASSERT_OK_AND_ASSIGN(std::vector<char> selected, Evaluate(result->RowFilter(), batch)); + ASSERT_EQ(std::vector<char>({0, 1, 0, 1}), selected); +} + +TEST(TableQueryAuthResultTest, NullLiteralsMatchAsInJava) { + auto batch = MakeBatch(R"([[1, "a", null, 1], [2, "b", null, 2], [3, null, null, 3]])"); + auto evaluate = [&](const std::string& filter) -> Result<std::vector<char>> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({filter}, {})); + return Evaluate(result->RowFilter(), batch); + }; + const std::vector<std::pair<std::string, std::vector<char>>> cases = { + {LeafJson(IdRef(), "EQUAL", "[null]"), {0, 0, 0}}, + {LeafJson(IdRef(), "IN", "[1, null]"), {1, 0, 0}}, + {LeafJson(IdRef(), "IN", "[null]"), {0, 0, 0}}, + {LeafJson(NameRef(), "IN", "[]"), {0, 0, 0}}, + {LeafJson(IdRef(), "NOT_IN", "[1, null]"), {0, 0, 0}}, + {LeafJson(IdRef(), "NOT_IN", "[1]"), {0, 1, 1}}, + {LeafJson(NameRef(), "NOT_IN", "[]"), {1, 1, 0}}, + {LeafJson(NameRef(), "NOT_IN", "[null]"), {0, 0, 0}}, + {LeafJson(IdRef(), "BETWEEN", "[2, 2]"), {0, 1, 0}}, + {LeafJson(IdRef(), "NOT_BETWEEN", "[2, 2]"), {1, 0, 1}}, + {LeafJson(NameRef(), "IS_NOT_NULL", "[]"), {1, 1, 0}}}; + for (const auto& [filter, expected] : cases) { + ASSERT_OK_AND_ASSIGN(std::vector<char> selected, evaluate(filter)); + ASSERT_EQ(expected, selected) << filter; + } +} + +TEST(TableQueryAuthResultTest, FoldsConstantFiltersAsJava) { + const std::string always_true = ConstantJson("TRUE"); + const std::string always_false = ConstantJson("FALSE"); + const std::string id_is_1 = LeafJson(IdRef(), "EQUAL", "[1]"); + + for (const std::string& filter : {always_true, LeafJson(NameRef(), "TRUE", "[]"), + CompoundJson("AND", always_true, always_true), + CompoundJson("OR", always_false, always_true), + CompoundJson("OR", id_is_1, always_true)}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({filter}, {})); + ASSERT_FALSE(result->HasRules()) << filter; + ASSERT_FALSE(result->RejectsAllRows()) << filter; + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> no_rules, + TableQueryAuthResult::Create({}, {})); + for (const std::vector<std::string>& filter : + std::vector<std::vector<std::string>>{{always_false}, + {LeafJson(NameRef(), "FALSE", "[]")}, + {CompoundJson("AND", id_is_1, always_false)}, + {CompoundJson("OR", always_false, always_false)}, + {id_is_1, always_false}}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create(filter, {})); + ASSERT_TRUE(result->HasRules()) << filter.back(); + ASSERT_TRUE(result->RejectsAllRows()) << filter.back(); + ASSERT_EQ(nullptr, result->RowFilter()) << filter.back(); + ASSERT_FALSE(result->SameFilter(*no_rules)) << filter.back(); + } + + auto batch = MakeBatch(R"([[1, "a", null, 1], [2, "b", null, 2]])"); + for (const std::string& filter : + {CompoundJson("AND", id_is_1, always_true), CompoundJson("OR", always_false, id_is_1)}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({filter}, {})); + ASSERT_FALSE(result->RejectsAllRows()) << filter; + ASSERT_EQ(std::set<std::string>({"id"}), result->RowFilterFields()) << filter; + ASSERT_OK_AND_ASSIGN(std::vector<char> selected, Evaluate(result->RowFilter(), batch)); + ASSERT_EQ(std::vector<char>({1, 0}), selected) << filter; + } + + Status unsupported_beside_constant = + TableQueryAuthResult::Create( + {CompoundJson("OR", always_true, LeafJson(NameRef(), "NOT_LIKE", R"(["a%"])"))}, {}) + .status(); + ASSERT_TRUE(unsupported_beside_constant.IsNotImplemented()) + << unsupported_beside_constant.ToString(); + Status computed_constant = + TableQueryAuthResult::Create( + {R"({"kind": "LEAF", "transform": {"name": "NULL"}, "function": "EQUAL",)" + R"( "literals": [1]})"}, + {}) + .status(); + ASSERT_TRUE(computed_constant.IsNotImplemented()) << computed_constant.ToString(); + + for (const std::string function : {"TRUE", "FALSE"}) { + for (const std::string literals : {"{}", R"("x")", "1"}) { + for (const std::string& filter : + {ConstantJson(function, literals), LeafJson(NameRef(), function, literals)}) { + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({filter}, {}), + "The 'literals' of a row filter must be a JSON array."); + } + } + } +} + +TEST(TableQueryAuthResultTest, MatchesLikeAsJava) { + auto batch = MakeBatch( + "[[1, \"a\\nb\", null, 1], [2, \"a\\r\\nb\", null, 2], [3, \"a\u2028b\", null, 3]," + " [4, \"a_b\", null, 4], [5, \"a%b\", null, 5], [6, \"\", null, 6]," + " [7, null, null, 7], [8, \"名字\", null, 8]]"); + const std::vector<std::pair<std::string, std::vector<char>>> cases = { + {"a_b", {1, 0, 1, 1, 1, 0, 0, 0}}, {"a__b", {0, 1, 0, 0, 0, 0, 0, 0}}, + {"a%b", {1, 1, 1, 1, 1, 0, 0, 0}}, {R"(a\\_b)", {0, 0, 0, 1, 0, 0, 0, 0}}, + {R"(a\\%b)", {0, 0, 0, 0, 1, 0, 0, 0}}, {"", {0, 0, 0, 0, 0, 1, 0, 0}}, + {"%", {1, 1, 1, 1, 1, 1, 0, 1}}, {"名_", {0, 0, 0, 0, 0, 0, 0, 1}}}; + for (const auto& [pattern, expected] : cases) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {LeafJson(NameRef(), "LIKE", fmt::format(R"(["{}"])", pattern))}, {})); + ASSERT_OK_AND_ASSIGN(std::vector<char> selected, Evaluate(result->RowFilter(), batch)); + ASSERT_EQ(expected, selected) << pattern; + } +} + +TEST(TableQueryAuthResultTest, ParsesLiteralsOfSupportedTypes) { + auto literal = [](const std::string& field_ref, const std::string& literals) { + return FilterLiteral(field_ref, "EQUAL", literals); + }; + ASSERT_OK_AND_ASSIGN(Literal boolean, literal(FieldRefJson(0, "f", "BOOLEAN"), "[true]")); + ASSERT_TRUE(boolean.GetValue<bool>()); + ASSERT_OK_AND_ASSIGN(Literal tinyint, literal(FieldRefJson(0, "f", "TINYINT"), "[-3]")); + ASSERT_EQ(-3, tinyint.GetValue<int8_t>()); + ASSERT_OK_AND_ASSIGN(Literal smallint, literal(FieldRefJson(0, "f", "SMALLINT"), "[300]")); + ASSERT_EQ(300, smallint.GetValue<int16_t>()); + ASSERT_OK_AND_ASSIGN(Literal bigint, literal(FieldRefJson(0, "f", "BIGINT"), "[5000000000]")); + ASSERT_EQ(5000000000LL, bigint.GetValue<int64_t>()); + ASSERT_OK_AND_ASSIGN(Literal float_literal, literal(FieldRefJson(0, "f", "FLOAT"), "[0.25]")); + ASSERT_FLOAT_EQ(0.25f, float_literal.GetValue<float>()); + ASSERT_OK_AND_ASSIGN(Literal double_literal, literal(FieldRefJson(0, "f", "DOUBLE"), "[1.5]")); + ASSERT_DOUBLE_EQ(1.5, double_literal.GetValue<double>()); + ASSERT_OK_AND_ASSIGN(Literal varchar, literal(FieldRefJson(0, "f", "VARCHAR(10)"), R"(["x"])")); + ASSERT_EQ("x", varchar.GetValue<std::string>()); + ASSERT_OK_AND_ASSIGN(Literal decimal_literal, + literal(FieldRefJson(0, "f", "DECIMAL(10, 2)"), R"(["12.3"])")); + auto decimal = decimal_literal.GetValue<Decimal>(); + ASSERT_EQ(10, decimal.Precision()); + ASSERT_EQ(2, decimal.Scale()); + ASSERT_EQ(1230, decimal.ToUnscaledLong()); + + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create( + {LeafJson(FieldRefJson(0, "f", "INT"), "EQUAL", R"(["1"])")}, {}), + "does not fit column 'f' of type int32"); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create( + {LeafJson(FieldRefJson(0, "f", "TINYINT"), "EQUAL", "[300]")}, {}), + "does not fit column 'f'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(FieldRefJson(0, "f", "DECIMAL(3, 2)"), "EQUAL", R"(["123.45"])")}, {}), + "does not fit column 'f'"); +} + +TEST(TableQueryAuthResultTest, ParsesFloatLiteralsAsJava) { + const std::string field = FieldRefJson(0, "f", "FLOAT"); + for (const std::string value : + {"1e300", "-1e300", "3.403e38", "-3.403e38", R"("NaN")", R"("Infinity")"}) { + ASSERT_NOK_WITH_MSG(FilterLiteral(field, "EQUAL", "[" + value + "]"), + "does not fit column 'f' of type float"); + } + ASSERT_OK_AND_ASSIGN(Literal maximum, FilterLiteral(field, "EQUAL", "[3.4028234663852886e38]")); + ASSERT_EQ(std::numeric_limits<float>::max(), maximum.GetValue<float>()); + ASSERT_OK_AND_ASSIGN(Literal minimum, + FilterLiteral(field, "EQUAL", "[-3.4028234663852886e38]")); + ASSERT_EQ(std::numeric_limits<float>::lowest(), minimum.GetValue<float>()); + ASSERT_OK_AND_ASSIGN(Literal java_maximum, FilterLiteral(field, "EQUAL", "[3.4028235e38]")); + ASSERT_EQ(std::numeric_limits<float>::max(), java_maximum.GetValue<float>()); + ASSERT_OK_AND_ASSIGN(Literal java_minimum, FilterLiteral(field, "EQUAL", "[-3.4028235e38]")); + ASSERT_EQ(std::numeric_limits<float>::lowest(), java_minimum.GetValue<float>()); + ASSERT_OK_AND_ASSIGN(Literal double_literal, + FilterLiteral(FieldRefJson(0, "f", "DOUBLE"), "EQUAL", "[1e300]")); + ASSERT_DOUBLE_EQ(1e300, double_literal.GetValue<double>()); + + const float positive = std::ldexp(1.0f, 60) + std::ldexp(1.0f, 37); + const float unsigned_value = std::ldexp(1.0f, 63) + std::ldexp(1.0f, 40); + const float big_integer = std::ldexp(1.0f, 64) + std::ldexp(1.0f, 41); + for (const auto& [text, expected] : std::vector<std::pair<std::string, float>>{ + {"1152921573326323713", positive}, + {"-1152921573326323713", -positive}, + {"9223372586610589697", unsigned_value}, + {"-9223372586610589697", -unsigned_value}, + {"18446745173221179393", big_integer}, + {"-18446745173221179393", -big_integer}, + {"340282346638528859811704183484516925440", std::numeric_limits<float>::max()}}) { + ASSERT_OK_AND_ASSIGN(Literal literal, FilterLiteral(field, "EQUAL", "[" + text + "]")); + ASSERT_EQ(expected, literal.GetValue<float>()) << text; + } + for (const std::string text : {"18446745173221179393.0", "1.8446745173221179393e19"}) { + ASSERT_OK_AND_ASSIGN(Literal literal, FilterLiteral(field, "EQUAL", "[" + text + "]")); + ASSERT_EQ(std::ldexp(1.0f, 64), literal.GetValue<float>()) << text; + } + ASSERT_NOK_WITH_MSG(FilterLiteral(field, "EQUAL", "[340282356779733661637539395458142568448]"), + "does not fit column 'f' of type float"); +} + +TEST(TableQueryAuthResultTest, PreservesDoubleLiteralPrecisionAsJava) { + const std::string field = FieldRefJson(0, "f", "DOUBLE"); + const double infinity = std::numeric_limits<double>::infinity(); + for (const auto& [text, expected] : std::vector<std::pair<std::string, double>>{ + {"0.94109657376859879", 0.94109657376859879}, + {"0.92057493778151411", 0.92057493778151411}, + {"-0.94109657376859879", -0.94109657376859879}, + {"0.986801886310680574", 0.986801886310680574}, + {"1.7976931348623157e308", std::numeric_limits<double>::max()}, + {"4.9406564584124654e-324", std::numeric_limits<double>::denorm_min()}}) { + SCOPED_TRACE(text); + ASSERT_OK_AND_ASSIGN(Literal literal, FilterLiteral(field, "EQUAL", "[" + text + "]")); + ASSERT_EQ(expected, literal.GetValue<double>()); + + arrow::DoubleBuilder builder; + ASSERT_TRUE(builder.Append(std::nextafter(expected, -infinity)).ok()); + ASSERT_TRUE(builder.Append(expected).ok()); + ASSERT_TRUE(builder.Append(std::nextafter(expected, infinity)).ok()); + ASSERT_TRUE(builder.AppendNull().ok()); + std::shared_ptr<arrow::Array> values; + ASSERT_TRUE(builder.Finish(&values).ok()); + auto batch = + arrow::StructArray::Make({values}, {arrow::field("f", arrow::float64())}).ValueOrDie(); + for (const auto& [function, selected] : + std::vector<std::pair<std::string, std::vector<char>>>{ + {"EQUAL", {0, 1, 0, 0}}, {"GREATER_THAN", {0, 0, 1, 0}}}) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> rules, + TableQueryAuthResult::Create({LeafJson(field, function, "[" + text + "]")}, {})); + ASSERT_OK_AND_ASSIGN(std::vector<char> actual, Evaluate(rules->RowFilter(), batch)); + ASSERT_EQ(selected, actual) << function; + } + } +} + +TEST(TableQueryAuthResultTest, RefusesMalformedRuleJson) { + const std::string suffix = std::string(1, '\0') + "not JSON"; + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({ConstantJson("TRUE") + suffix}, {}), + "Row filter contains an unescaped NUL byte"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create({}, {{"name", R"({"name": "NULL"})" + suffix}}), + "Column mask transform contains an unescaped NUL byte"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "CONCAT", "inputs": ["a\u0000b"]})"}})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> masked, + Mask(*result, "name", MakeBatch(R"([[1, "x", null, 1]])"), arrow::utf8())); + ASSERT_TRUE(masked->Equals(*Strings(R"(["a\u0000b"])"))); + + for (const std::string& invalid : + {std::string(1, '\xFF'), std::string("\xC3\x28", 2), std::string("\xED\xA0\x80", 3)}) { + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create( + {LeafJson(NameRef(), "EQUAL", "[\"" + invalid + "\"]")}, {}), + "Row filter is not valid JSON"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "CONCAT", "inputs": [")" + invalid + R"("]})"}}), + "Column mask transform is not valid JSON"); + } + + ASSERT_NOK(TableQueryAuthResult::Create( + {LeafJson(R"({"index": 0, "name": "id", "type": "INT\u0000suffix"})", "EQUAL", "[1]")}, + {})); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(FieldRefJson(0, "f", "DECIMAL(39, 0)"), "EQUAL", R"(["1"])")}, {}), + "Decimal precision out of range"); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create( + {}, {{"name", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", + FieldRefJson(1, "name", "DECIMAL(39, 0)"))}}), + "Decimal precision out of range"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "FIELD_REF", "fieldRef": {"index": 1, "name": "name",)" + R"("type": {"type": "ARRAY\u0000suffix", "element": "STRING"}}})"}}), + "complex data type cannot contain a NUL byte"); + + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({"{"}, {}), "Row filter is not valid JSON"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {R"({"kind": "LEAF", "transform": {"name": "NULL"}, "function": "TRUE",)" + R"( "function": "FALSE", "literals": []})"}, + {}), + "Row filter is ambiguous: JSON object names member 'function' more than once."); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {fmt::format(R"({{"kind": "COMPOUND", "function": "AND", "children": [{}],)" + R"( "children": [{}]}})", + ConstantJson("TRUE"), ConstantJson("FALSE"))}, + {}), + "member 'children'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(R"({"index": 1, "name": "name", "name": "phone", "type": "STRING"})", + "IS_NULL", "[]")}, + {}), + "member 'name'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(R"({"index": 0, "name": "f", "type": {"type": "ARRAY", "element": "INT",)" + R"( "element": "STRING"}})", + "IS_NULL", "[]")}, + {}), + "member 'element'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}, "name": "NULL"}})", + NameRef())}}), + "Column mask transform is ambiguous: JSON object names member 'name' more than once."); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "CONCAT", "inputs": ["***"], "inputs": []})"}}), + "member 'inputs'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "SUBSTRING", "inputs": [{"index": 1, "name": "name",)" + R"( "type": "STRING", "type": "INT"}, 1, 3]})"}}), + "member 'type'"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", R"({"name": "NULL", "n\u0061me": "CONCAT", "inputs": ["***"]})"}}), + "member 'name'"); +} + +TEST(TableQueryAuthResultTest, ParsesTemporalLiteralsAsJava) { + auto date = [](const std::string& literal) -> Result<int32_t> { + PAIMON_ASSIGN_OR_RAISE(Literal parsed, FilterLiteral(FieldRefJson(0, "d", "DATE"), "EQUAL", + fmt::format("[{}]", literal))); + return parsed.GetValue<int32_t>(); + }; + for (const auto& [text, epoch_day] : + std::vector<std::pair<std::string, int32_t>>{{"1970-01-02", 1}, + {"1969-12-31", -1}, + {"2024-02-29", 19782}, + {"-0001-01-01", -719893}, + {"+10000-01-01", 2932897}}) { + ASSERT_OK_AND_ASSIGN(int32_t parsed, date(fmt::format(R"("{}")", text))); + ASSERT_EQ(epoch_day, parsed) << text; + } + for (const std::string text : + {"123", "2024-01-01junk", "2024-01-01 ", " 2024-01-01", "2023-02-29", "2024-13-01", + "2024-1-01", "+2024-01-01", "10000-01-01", "-0000-01-01"}) { + ASSERT_NOK_WITH_MSG(date(fmt::format(R"("{}")", text)), + fmt::format("Cannot parse '{}' as a date literal", text)); + } + ASSERT_NOK_WITH_MSG(date("19782"), "does not fit column 'd' of type date32"); + + auto timestamp = [](const std::string& type, const std::string& text) -> Result<Timestamp> { + PAIMON_ASSIGN_OR_RAISE(Literal literal, + FilterLiteral(FieldRefJson(0, "ts", type), "GREATER_THAN", + fmt::format(R"(["{}"])", text))); + return literal.GetValue<Timestamp>(); + }; + const std::string local = "TIMESTAMP(9)"; + const std::string instant = "TIMESTAMP(9) WITH LOCAL TIME ZONE"; + for (const auto& [type, text, millis, nano_of_millis] : + std::vector<std::tuple<std::string, std::string, int64_t, int32_t>>{ + {local, "2263-01-01T00:00:00", 9246182400000LL, 0}, + {local, "1600-01-01T00:00:00", -11676096000000LL, 0}, + {local, "1969-12-31T23:59:59.5", -500, 0}, + {local, "1970-01-01T00:00:00.123456789", 123, 456789}, + {local, "1970-01-01T00:01", 60000, 0}, + {"TIMESTAMP(6)", "1970-01-01T00:00:01.000001", 1000, 1000}, + {local, "9999-12-31T23:59:59.999999999", 253402300799999LL, 999999}, + {local, "+10000-01-01T00:00", 253402300800000LL, 0}, + {local, "+292278994-08-17T07:12:55.807999999", std::numeric_limits<int64_t>::max(), + 999999}, + {local, "-292275055-05-16T16:47:04.192", std::numeric_limits<int64_t>::min(), 0}, + {instant, "1969-12-31T23:59:59.5Z", -500, 0}, + {instant, "+10000-01-01T00:00:00Z", 253402300800000LL, 0}, + {instant, "+292278994-08-17T07:12:55.807Z", std::numeric_limits<int64_t>::max(), 0}, + {instant, "-292275055-05-16T16:47:04.192999999Z", std::numeric_limits<int64_t>::min(), + 999999}}) { + ASSERT_OK_AND_ASSIGN(Timestamp parsed, timestamp(type, text)); + ASSERT_EQ(millis, parsed.GetMillisecond()) << text; + ASSERT_EQ(nano_of_millis, parsed.GetNanoOfMillisecond()) << text; + } + for (const auto& [type, text] : std::vector<std::pair<std::string, std::string>>{ + {local, "1970-01-01.5"}, + {local, "1970-01-01T00:00.5"}, + {local, "1970-01-01"}, + {local, "1970-01-01T24:00:00"}, + {local, "1970-01-01T00:00:00Z"}, + {local, "1970-01-01T00:00:00."}, + {local, "1970-01-01T00:00:00.1234567890"}, + {local, "1970-01-01T00:00:00." + std::string(64, '9')}, + {instant, "1970-01-01T00:00:01"}, + {instant, "1970-01-01T00:01Z"}, + {instant, "1970-01-01T00:00:01+08:00"}}) { + ASSERT_NOK_WITH_MSG(timestamp(type, text), + fmt::format("Cannot parse '{}' as a timestamp literal", text)); + } + ASSERT_NOK_WITH_MSG(timestamp(local, "+292278995-01-01T00:00:00"), + "is outside the range of epoch milliseconds"); + for (const std::string& type : {local, instant}) { + const std::string suffix = type == instant ? "Z" : ""; + for (const std::string text : + {"+292278994-08-17T07:12:55.808", "-292275055-05-16T16:47:04.191999999"}) { + ASSERT_NOK_WITH_MSG(timestamp(type, text + suffix), + "is outside the range of epoch milliseconds"); + } + } + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> rule, + TableQueryAuthResult::Create({LeafJson(FieldRefJson(0, "ts", "TIMESTAMP(3)"), + "GREATER_THAN", R"(["2263-01-01T00:00:00"])")}, + {})); + auto batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MILLI))}), + "[[1767225600000], [9246182400001]]") + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::vector<char> selected, Evaluate(rule->RowFilter(), batch)); + ASSERT_EQ(std::vector<char>({0, 1}), selected); + + // A cutoff beyond epoch nanoseconds must not wrap when compared with a TIMESTAMP(9) column. + const std::string cutoff = "2500-01-01T00:00:00.000000001"; + auto nanos_batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::NANO)), + arrow::field("allowed", arrow::boolean())}), + "[[1577836800000000000, true], [1577836800000000000, false]]") + .ValueOrDie()); + for (const char* function : {"GREATER_THAN", "GREATER_OR_EQUAL"}) { + SCOPED_TRACE(function); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> cutoff_rule, + TableQueryAuthResult::Create( + {CompoundJson("OR", + LeafJson(FieldRefJson(0, "ts", local), function, + fmt::format(R"(["{}"])", cutoff)), + LeafJson(FieldRefJson(1, "allowed", "BOOLEAN"), "EQUAL", "[true]"))}, + {})); + ASSERT_OK_AND_ASSIGN(selected, Evaluate(cutoff_rule->RowFilter(), nanos_batch)); + ASSERT_EQ(std::vector<char>({1, 0}), selected); + } + for (const char* function : {"LESS_THAN", "NOT_EQUAL"}) { + SCOPED_TRACE(function); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> cutoff_rule, + TableQueryAuthResult::Create({LeafJson(FieldRefJson(0, "ts", local), function, + fmt::format(R"(["{}"])", cutoff))}, + {})); + ASSERT_OK_AND_ASSIGN(selected, Evaluate(cutoff_rule->RowFilter(), nanos_batch)); + ASSERT_EQ(std::vector<char>({1, 1}), selected); + } + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> in_rule, + TableQueryAuthResult::Create( + {CompoundJson("AND", + LeafJson(FieldRefJson(0, "ts", local), "IN", + fmt::format(R"(["{}", "2020-01-01T00:00:00"])", cutoff)), + LeafJson(FieldRefJson(1, "allowed", "BOOLEAN"), "EQUAL", "[true]"))}, + {})); + ASSERT_OK_AND_ASSIGN(selected, Evaluate(in_rule->RowFilter(), nanos_batch)); + ASSERT_EQ(std::vector<char>({1, 0}), selected); +} + +TEST(TableQueryAuthResultTest, RefusesWhatItCannotApply) { + Status unknown_function = + TableQueryAuthResult::Create({LeafJson(NameRef(), "NOT_LIKE", R"(["a%"])")}, {}).status(); + ASSERT_TRUE(unknown_function.IsNotImplemented()) << unknown_function.ToString(); + Status unknown_kind = + TableQueryAuthResult::Create({R"({"kind": "OTHER", "function": "AND"})"}, {}).status(); + ASSERT_TRUE(unknown_kind.IsNotImplemented()) << unknown_kind.ToString(); + Status computed_filter = + TableQueryAuthResult::Create( + {R"({"kind": "LEAF", "transform": {"name": "UPPER", "inputs": []},)" + R"( "function": "EQUAL", "literals": ["A"]})"}, + {}) + .status(); + ASSERT_TRUE(computed_filter.IsNotImplemented()) << computed_filter.ToString(); + Status unknown_mask = + TableQueryAuthResult::Create( + {}, {{"name", fmt::format(R"({{"name": "UPPER", "inputs": [{}]}})", NameRef())}}) + .status(); + ASSERT_TRUE(unknown_mask.IsNotImplemented()) << unknown_mask.ToString(); + + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({""}, {}), "Row filter cannot be empty."); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({"null"}, {}), + "Row filter cannot be JSON null."); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({}, {{"", R"({"name": "NULL"})"}}), + "Column mask target cannot be empty."); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({}, {{"name", ""}}), + "Column mask transform cannot be empty."); + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({}, {{"name", "null"}}), + "Column mask transform cannot be JSON null."); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create({LeafJson(FieldRefJson(0, "id", "INT"), "IS_NULL", "[]"), + LeafJson(FieldRefJson(0, "id", "BIGINT"), "IS_NULL", "[]")}, + {}), + "read column 'id' as both"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(R"({"index": 0, "name": "f", "type": {"type": 123}})", "IS_NULL", "[]")}, {}), + "complex data type must have a string type"); + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {LeafJson(R"({"index": 0, "name": "f", "type": {"type": "ARRAY", "element": 5}})", + "IS_NULL", "[]")}, + {}), + "cannot parse data type"); + for (const char* inputs : {"[]", R"(["-"])"}) { + ASSERT_NOK_WITH_MSG( + TableQueryAuthResult::Create( + {}, {{"name", fmt::format(R"({{"name": "CONCAT_WS", "inputs": {}}})", inputs)}}), + "A CONCAT_WS transform needs a separator and at least one input."); + } + for (const std::string function : {"IN", "NOT_IN"}) { + const std::string missing_literals = + fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "FIELD_REF", "fieldRef": {}}},)" + R"( "function": "{}"}})", + NameRef(), function); + for (const std::string& filter : + {missing_literals, LeafJson(NameRef(), function, "null")}) { + ASSERT_NOK_WITH_MSG(TableQueryAuthResult::Create({filter}, {}), + fmt::format("Invalid: Query auth row filter {} requires a " + "non-null 'literals' array.", + function)); + } + } +} + +TEST(TableQueryAuthResultTest, AppliesColumnMasks) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {}, + {{"salary", R"({"name": "NULL"})"}, + {"name", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", PhoneRef())}, + {"phone", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, 1, 3]}})", PhoneRef())}, + {"concat", fmt::format(R"({{"name": "CONCAT", "inputs": [{}, "-", {}]}})", NameRef(), + PhoneRef())}, + {"concat_ws", fmt::format(R"({{"name": "CONCAT_WS", "inputs": ["-", {}, null, )" + R"("x", {}]}})", + NameRef(), PhoneRef())}, + {"null_separator", + fmt::format(R"({{"name": "CONCAT_WS", "inputs": [null, {}]}})", NameRef())}, + {"empty", R"({"name": "CONCAT_WS", "inputs": ["-", ""]})"}, + {"tail", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, -4]}})", PhoneRef())}, + {"utf8", R"({"name": "SUBSTRING", "inputs": ["héllo", 2, 3]})"}, + {"null_position", + fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, null]}})", PhoneRef())}})); + ASSERT_TRUE(result->HasRules()); + ASSERT_EQ(nullptr, result->RowFilter()); + auto batch = MakeBatch(R"([[1, "bob", "13900002222", 10], [2, null, null, 20]])"); + + auto assert_mask = [&](const std::string& column, const std::shared_ptr<arrow::DataType>& type, + const std::shared_ptr<arrow::Array>& expected) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> masked, + Mask(*result, column, batch, type)); + ASSERT_TRUE(masked->Equals(*expected)) + << column << ": " << masked->ToString() << " vs " << expected->ToString(); + }; + assert_mask( + "salary", arrow::int64(), + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[null, null]").ValueOrDie()); + assert_mask("name", arrow::utf8(), Strings(R"(["13900002222", null])")); + assert_mask("phone", arrow::utf8(), Strings(R"(["139", null])")); + assert_mask("concat", arrow::utf8(), Strings(R"(["bob-13900002222", null])")); + assert_mask("concat_ws", arrow::utf8(), Strings(R"(["bob-x-13900002222", "x"])")); + assert_mask("null_separator", arrow::utf8(), Strings("[null, null]")); + assert_mask("empty", arrow::utf8(), Strings(R"(["", ""])")); + assert_mask("tail", arrow::utf8(), Strings(R"(["2222", null])")); + assert_mask("utf8", arrow::utf8(), Strings("[\"\xC3\xA9ll\", \"\xC3\xA9ll\"]")); + assert_mask("null_position", arrow::utf8(), Strings("[null, null]")); +} + +TEST(TableQueryAuthResultTest, SubstringReadsPositionsAsJava) { + const auto type = + arrow::struct_({arrow::field("phone", arrow::utf8()), arrow::field("pos", arrow::int64()), + arrow::field("len", arrow::int64())}); + auto batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON(type, R"([["abc", 9223372036854775807, null], + ["abc", null, 9223372036854775807], + [null, 9223372036854775807, 9223372036854775807], + ["abc", 2, 1]])") + .ValueOrDie()); + for (const auto& [pos, len, expected] : + std::vector<std::tuple<std::string, std::string, std::string>>{ + {FieldRefJson(1, "pos", "BIGINT"), FieldRefJson(2, "len", "BIGINT"), + R"([null, null, null, "b"])"}, + {"null", FieldRefJson(1, "pos", "BIGINT"), "[null, null, null, null]"}, + {FieldRefJson(2, "len", "BIGINT"), "null", "[null, null, null, null]"}}) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {}, {{"phone", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, {}, {}]}})", + FieldRefJson(0, "phone", "STRING"), pos, len)}})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> masked, + Mask(*result, "phone", batch, arrow::utf8())); + ASSERT_TRUE(masked->Equals(*Strings(expected))) << masked->ToString(); + } + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {}, + {{"phone", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, {}, 3]}})", PhoneRef(), + FieldRefJson(1, "secret_position", "BIGINT"))}})); + auto secret_batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("phone", arrow::utf8()), + arrow::field("secret_position", arrow::int64())}), + R"([["13900001111", 9223372036854775807]])") + .ValueOrDie()); + const Status status = Mask(*result, "phone", secret_batch, arrow::utf8()).status(); + ASSERT_TRUE(status.IsInvalid()) << status.ToString(); + ASSERT_NOK_WITH_MSG(status, "SUBSTRING position of column 'secret_position' does not fit"); + ASSERT_EQ(std::string::npos, status.message().find("9223372036854775807")); +} + +TEST(TableQueryAuthResultTest, CopiesATimeColumnThatNoFilterCanCompare) { + const std::string src_ref = FieldRefJson(0, "src", "TIME(3)"); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {}, {{"dst", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", src_ref)}})); + const auto time_type = arrow::time32(arrow::TimeUnit::MILLI); + ASSERT_OK(result->ValidateAgainstSchema({DataField(0, arrow::field("src", time_type)), + DataField(1, arrow::field("dst", time_type))})); + ASSERT_NOK_WITH_MSG( + result->ValidateAgainstSchema({DataField(0, arrow::field("src", time_type)), + DataField(1, arrow::field("dst", arrow::int32()))}), + "copies column 'src' of type time32[ms] into a column of type int32"); + auto batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("src", time_type), arrow::field("dst", time_type)}), + "[[1000, 5], [null, 6]]") + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> masked, + Mask(*result, "dst", batch, time_type)); + ASSERT_TRUE(masked->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(time_type, "[1000, null]").ValueOrDie())) + << masked->ToString(); + + Status filter = TableQueryAuthResult::Create({LeafJson(src_ref, "IS_NULL", "[]")}, {}).status(); + ASSERT_TRUE(filter.IsNotImplemented()) << filter.ToString(); + ASSERT_NOK_WITH_MSG(filter, "Query auth row filter IS_NULL on column 'src' of type time32[ms]"); +} + +TEST(TableQueryAuthResultTest, ValidatesTheRulesAgainstTheSchema) { + const std::vector<DataField> fields = TableFields(); + auto validate = [&](const std::vector<std::string>& filter, + const std::map<std::string, std::string>& masks) -> Status { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create(filter, masks)); + return result->ValidateAgainstSchema(fields); + }; + ASSERT_OK(validate( + {LeafJson(IdRef(), "GREATER_THAN", "[1]")}, + {{"salary", R"({"name": "NULL"})"}, + {"phone", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, 1, 3]}})", PhoneRef())}})); + + ASSERT_NOK_WITH_MSG(validate({LeafJson(FieldRefJson(9, "gone", "INT"), "IS_NULL", "[]")}, {}), + "Row filter references column 'gone' which does not exist"); + const std::string gone_is_1 = LeafJson(FieldRefJson(9, "gone", "INT"), "EQUAL", "[1]"); + for (const std::string& folded : {CompoundJson("OR", ConstantJson("TRUE"), gone_is_1), + CompoundJson("AND", ConstantJson("FALSE"), gone_is_1), + LeafJson(FieldRefJson(9, "gone", "INT"), "TRUE", "[]")}) { + ASSERT_NOK_WITH_MSG(validate({folded}, {}), + "A query auth rule references column 'gone' which does not exist"); + } + ASSERT_NOK_WITH_MSG(validate({}, {{"gone", R"({"name": "NULL"})"}}), + "Column masking references column 'gone' which does not exist"); + ASSERT_NOK_WITH_MSG(validate({LeafJson(FieldRefJson(3, "salary", "INT"), "IS_NULL", "[]")}, {}), + "read column 'salary' as int32, but the table holds it as int64"); + ASSERT_NOK_WITH_MSG( + validate({}, + {{"salary", fmt::format(R"({{"name": "CONCAT", "inputs": [{}]}})", NameRef())}}), + "produces a string, but column 'salary' is int64"); + ASSERT_NOK_WITH_MSG( + validate({}, + {{"salary", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", IdRef())}}), + "copies column 'id' of type int32 into a column of type int64"); + ASSERT_NOK_WITH_MSG( + validate({}, + {{"name", R"({"name": "NULL"})"}, + {"phone", fmt::format(R"({{"name": "CONCAT", "inputs": [{}]}})", NameRef())}}), + "Column masking on 'phone' reads column 'name', which is masked too"); +} + +TEST(TableQueryAuthResultTest, LeavesMasksOfSystemColumnsInert) { + const std::string row_id_ref = FieldRefJson(-1, "_ROW_ID", "BIGINT NOT NULL"); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> system_mask, + TableQueryAuthResult::Create( + {}, {{"_ROW_ID", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", row_id_ref)}, + {"rowkind", R"({"name": "NULL"})"}, + {"name", R"({"name": "NULL"})"}})); + ASSERT_EQ(std::set<std::string>{"name"}, system_mask->MaskedFields()); + ASSERT_EQ(std::set<std::string>{"name"}, system_mask->ReferencedFields()); + ASSERT_EQ(3u, system_mask->ColumnMasking().size()); + ASSERT_OK(system_mask->ValidateAgainstSchema(TableFields())); + Status unknown_system_mask = + TableQueryAuthResult::Create({}, {{"_ROW_ID", R"({"name": "UPPER", "inputs": []})"}}) + .status(); + ASSERT_TRUE(unknown_system_mask.IsNotImplemented()) << unknown_system_mask.ToString(); + + for (const std::string name : + {"_SEQUENCE_NUMBER", "_VALUE_KIND", "_LEVEL", "rowkind", "_ROW_ID"}) { + SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> rules, + TableQueryAuthResult::Create({}, {{name, R"({"name": "NULL"})"}})); + ASSERT_OK(rules->ValidateAgainstSchema(TableFields())); + ASSERT_NOK_WITH_MSG( + rules->ValidateAgainstSchema({DataField(0, arrow::field(name, arrow::utf8()))}), + "refusing to ignore its mask"); + } +} + +TEST(TableQueryAuthResultTest, RefusesAColumnRenamedOrReplacedSinceTheRulesWereBound) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({LeafJson(NameRef(), "IS_NULL", "[]")}, + {{"phone", R"({"name": "NULL"})"}})); + const std::vector<DataField> latest = TableFields(); + const std::map<std::string, int32_t> field_ids = result->BindFieldIds(latest); + ASSERT_EQ((std::map<std::string, int32_t>{{"name", 1}, {"phone", 2}}), field_ids); + ASSERT_OK(result->ValidateFieldIds(field_ids, latest)); + + std::vector<DataField> renamed = latest; + renamed[1] = DataField(1, arrow::field("old_name", arrow::utf8())); + ASSERT_NOK_WITH_MSG(result->ValidateFieldIds(field_ids, renamed), + "exposes as 'old_name' (renamed since)"); + + std::vector<DataField> replaced = latest; + replaced[1] = DataField(7, arrow::field("name", arrow::utf8())); + ASSERT_NOK_WITH_MSG(result->ValidateFieldIds(field_ids, replaced), + "as a different column of the same name"); + + std::vector<DataField> swapped = latest; + swapped[1] = DataField(1, arrow::field("phone", arrow::utf8())); + swapped[2] = DataField(2, arrow::field("name", arrow::utf8())); + ASSERT_OK(result->ValidateAgainstSchema(swapped)); + ASSERT_NOK_WITH_MSG(result->ValidateFieldIds(field_ids, swapped), + "as a different column of the same name"); + + std::map<std::string, int32_t> incomplete = field_ids; + incomplete.erase("phone"); + ASSERT_NOK_WITH_MSG(result->ValidateFieldIds(incomplete, latest), + "Query auth rules reference column 'phone', which the schema the scan " + "planned with did not hold"); + ASSERT_EQ((std::set<std::string>{"name", "phone"}), result->ReferencedFields()); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> folded, + TableQueryAuthResult::Create( + {CompoundJson("OR", ConstantJson("TRUE"), LeafJson(NameRef(), "IS_NULL", "[]"))}, {})); + ASSERT_EQ(nullptr, folded->RowFilter()); + ASSERT_EQ((std::set<std::string>{"name"}), folded->ReferencedFields()); + const std::map<std::string, int32_t> folded_ids = folded->BindFieldIds(latest); + ASSERT_EQ((std::map<std::string, int32_t>{{"name", 1}}), folded_ids); + ASSERT_NOK_WITH_MSG(folded->ValidateFieldIds(folded_ids, renamed), + "exposes as 'old_name' (renamed since)"); + ASSERT_NOK_WITH_MSG(folded->ValidateFieldIds(folded_ids, replaced), + "as a different column of the same name"); +} + +TEST(TableQueryAuthResultTest, RequiredFieldsFollowTheMasksTransitively) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create( + {LeafJson(IdRef(), "GREATER_THAN", "[1]")}, + {{"phone", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", NameRef())}, + {"salary", R"({"name": "NULL"})"}})); + ASSERT_EQ(std::set<std::string>({"id", "name"}), result->RequiredFields({"phone"})); + ASSERT_EQ(std::set<std::string>({"id"}), result->RequiredFields({"salary"})); + ASSERT_EQ(std::set<std::string>({"phone", "salary"}), result->MaskedFields()); +} + +TEST(TableQueryAuthResultTest, SameFilterComparesTheParsedConjuncts) { + const std::string first = LeafJson(IdRef(), "GREATER_THAN", "[1]"); + const std::string second = LeafJson(NameRef(), "IS_NULL", "[]"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + TableQueryAuthResult::Create({first, second}, {})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> reordered, + TableQueryAuthResult::Create({second, first}, {})); + ASSERT_TRUE(result->SameFilter(*reordered)); + const std::string reserialized = + fmt::format(R"({{"literals":[1],"function":"GREATER_THAN","kind":"LEAF",)" + R"("transform":{{"fieldRef":{},"name":"FIELD_REF"}}}})", + IdRef()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> compound, + TableQueryAuthResult::Create({CompoundJson("AND", second, reserialized)}, + {{"name", R"({"name": "NULL"})"}})); + ASSERT_TRUE(result->SameFilter(*compound)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> fewer, + TableQueryAuthResult::Create({first}, {})); + ASSERT_FALSE(result->SameFilter(*fewer)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> other, + TableQueryAuthResult::Create({first, LeafJson(NameRef(), "IS_NOT_NULL", "[]")}, {})); + ASSERT_FALSE(result->SameFilter(*other)); +} + +} // namespace paimon::test
diff --git a/src/paimon/core/core_options.cpp b/src/paimon/core/core_options.cpp index 4ebe5cc..e5fa9db 100644 --- a/src/paimon/core/core_options.cpp +++ b/src/paimon/core/core_options.cpp
@@ -490,6 +490,7 @@ bool commit_discard_duplicate_files = false; bool rest_commit_lock_enabled = false; bool rest_commit_acquire_lock = false; + bool query_auth_enabled = false; bool dynamic_partition_overwrite = true; bool overwrite_upgrade = true; bool compaction_force_rewrite_all_files = false; @@ -726,6 +727,8 @@ // Parse rest.commit.acquire-lock - whether this writer acquires a commit lease PAIMON_RETURN_NOT_OK( parser.Parse<bool>(Options::REST_COMMIT_ACQUIRE_LOCK, &rest_commit_acquire_lock)); + // Parse query-auth.enabled - whether the catalog enforces row filters and column masks + PAIMON_RETURN_NOT_OK(parser.Parse<bool>(Options::QUERY_AUTH_ENABLED, &query_auth_enabled)); // Parse dynamic-partition-overwrite - whether overwrite only dynamic partitions // for partitioned table overwrite. PAIMON_RETURN_NOT_OK( @@ -1382,6 +1385,10 @@ return impl_->rest_commit_acquire_lock; } +bool CoreOptions::QueryAuthEnabled() const { + return impl_->query_auth_enabled; +} + bool CoreOptions::DynamicPartitionOverwrite() const { return impl_->dynamic_partition_overwrite; }
diff --git a/src/paimon/core/core_options.h b/src/paimon/core/core_options.h index bfc6ee2..616296d 100644 --- a/src/paimon/core/core_options.h +++ b/src/paimon/core/core_options.h
@@ -145,6 +145,7 @@ bool CommitDiscardDuplicateFiles() const; bool RestCommitLockEnabled() const; bool RestCommitAcquireLock() const; + bool QueryAuthEnabled() const; bool DynamicPartitionOverwrite() const; bool OverwriteUpgrade() const; int32_t GetCompactionMinFileNum() const;
diff --git a/src/paimon/core/core_options_test.cpp b/src/paimon/core/core_options_test.cpp index 1344d91..e5b44dc 100644 --- a/src/paimon/core/core_options_test.cpp +++ b/src/paimon/core/core_options_test.cpp
@@ -110,6 +110,7 @@ ASSERT_FALSE(core_options.CommitDiscardDuplicateFiles()); ASSERT_FALSE(core_options.RestCommitLockEnabled()); ASSERT_FALSE(core_options.RestCommitAcquireLock()); + ASSERT_FALSE(core_options.QueryAuthEnabled()); ExpireConfig expire_config = core_options.GetExpireConfig(); ASSERT_EQ(10, expire_config.GetSnapshotRetainMin()); ASSERT_EQ(std::numeric_limits<int32_t>::max(), expire_config.GetSnapshotRetainMax()); @@ -250,6 +251,7 @@ {Options::COMMIT_DISCARD_DUPLICATE_FILES, "true"}, {Options::REST_COMMIT_LOCK_ENABLED, "true"}, {Options::REST_COMMIT_ACQUIRE_LOCK, "true"}, + {Options::QUERY_AUTH_ENABLED, "true"}, {Options::DYNAMIC_PARTITION_OVERWRITE, "false"}, {Options::OVERWRITE_UPGRADE, "false"}, {Options::SCAN_SNAPSHOT_ID, "5"}, @@ -407,6 +409,7 @@ ASSERT_TRUE(core_options.CommitDiscardDuplicateFiles()); ASSERT_TRUE(core_options.RestCommitLockEnabled()); ASSERT_TRUE(core_options.RestCommitAcquireLock()); + ASSERT_TRUE(core_options.QueryAuthEnabled()); ASSERT_EQ(5, core_options.GetScanSnapshotId().value_or(-1)); ASSERT_EQ(7, core_options.GetScanManifestEntryCacheMaxSnapshots()); ASSERT_FALSE(core_options.ScanManifestEntryLazyDecodeEnabled());
diff --git a/src/paimon/core/global_index/global_index_scan.cpp b/src/paimon/core/global_index/global_index_scan.cpp index 457c3a1..bd2b075 100644 --- a/src/paimon/core/global_index/global_index_scan.cpp +++ b/src/paimon/core/global_index/global_index_scan.cpp
@@ -23,6 +23,7 @@ #include "paimon/core/global_index/global_index_scan_impl.h" #include "paimon/core/operation/file_store_scan.h" #include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/schema_validation.h" #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" @@ -81,6 +82,10 @@ std::shared_ptr<MemoryPool> pool = memory_pool ? memory_pool : GetDefaultPool(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> table_schema, LoadSchema(root_path, options, file_system)); + // An index search reveals which rows match, so it must not see the rows a row filter hides, + // and it ranks raw values, which a column mask invalidates. + PAIMON_RETURN_NOT_OK( + SchemaValidation::RejectQueryAuth(*table_schema, options, "Global index search")); PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, MergeOptions(table_schema, options, file_system)); std::shared_ptr<PredicateFilter> partition_filters; @@ -111,6 +116,8 @@ std::shared_ptr<MemoryPool> pool = memory_pool ? memory_pool : GetDefaultPool(); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> table_schema, LoadSchema(root_path, options, file_system)); + PAIMON_RETURN_NOT_OK( + SchemaValidation::RejectQueryAuth(*table_schema, options, "Global index search")); PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, MergeOptions(table_schema, options, file_system)); PAIMON_ASSIGN_OR_RAISE(Snapshot snapshot, LoadSnapshot(root_path, snapshot_id, core_options));
diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index f2bf754..21d4852 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp
@@ -26,6 +26,8 @@ #include "arrow/c/bridge.h" #include "paimon/catalog/catalog.h" #include "paimon/common/utils/path_util.h" +#include "paimon/core/catalog/catalog_utils.h" +#include "paimon/core/table/system/system_table.h" #include "paimon/core/utils/branch_manager.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" @@ -411,15 +413,36 @@ if (impl_->catalog_ == nullptr && impl_->identifier_) { return Status::Invalid("cannot read through a null catalog"); } + // Resolve a branch named by a system table path too, before requesting its credentials. + std::optional<SystemTablePath> system_table_path; + if (impl_->format_table_ == nullptr) { + PAIMON_ASSIGN_OR_RAISE(system_table_path, SystemTableLoader::TryParsePath(impl_->path_)); + } + std::map<std::string, std::string> branch_options = impl_->options_; + if (system_table_path && system_table_path->branch) { + PAIMON_ASSIGN_OR_RAISE(std::string path_branch, + BranchManager::ResolveBranch(impl_->identifier_, branch_options, + system_table_path->branch, "read")); + branch_options[Options::BRANCH] = path_branch; + } + PAIMON_ASSIGN_OR_RAISE( + std::string branch, + BranchManager::ResolveBranch(impl_->identifier_, branch_options, impl_->branch_, "read")); + if (impl_->catalog_ != nullptr && !BranchManager::IsMainBranch(branch)) { + PAIMON_RETURN_NOT_OK(BranchManager::CheckCatalogAddressableBranch(branch)); + } + if (impl_->identifier_) { + PAIMON_ASSIGN_OR_RAISE( + Identifier branch_identifier, + CatalogUtils::ResolvedBranchIdentifier(impl_->identifier_.value(), branch)); + impl_->identifier_.emplace(std::move(branch_identifier)); + } if (impl_->catalog_ != nullptr && impl_->specific_file_system_ == nullptr) { // A catalog issuing per-table temporary credentials only hands them out through // GetTableFileSystem, so the read uses the credentials of this table. PAIMON_ASSIGN_OR_RAISE(impl_->specific_file_system_, impl_->catalog_->GetTableFileSystem(impl_->identifier_.value())); } - PAIMON_ASSIGN_OR_RAISE( - std::string branch, - BranchManager::ResolveBranch(impl_->identifier_, impl_->options_, impl_->branch_, "read")); if (impl_->enable_prefetch_ && impl_->prefetch_max_parallel_num_ == 0) { return Status::Invalid("prefetch max parallel num should be greater than 0"); }
diff --git a/src/paimon/core/operation/read_context_test.cpp b/src/paimon/core/operation/read_context_test.cpp index 3d5040a..c13fba6 100644 --- a/src/paimon/core/operation/read_context_test.cpp +++ b/src/paimon/core/operation/read_context_test.cpp
@@ -19,6 +19,7 @@ #include "paimon/read_context.h" #include <utility> +#include <vector> #include "arrow/c/bridge.h" #include "arrow/type.h" @@ -64,6 +65,94 @@ ASSERT_EQ(catalog->TableFileSystemRequests().size(), 1U); } +TEST(ReadContextTest, TestWithCatalogResolvesTheBranch) { + auto catalog = std::make_shared<MockVersionManagedCatalog>(); + catalog->SetTableFileSystem(std::make_shared<MockFileSystem>()); + + for (const std::string branch : {"MAIN", "Main"}) { + const std::string message = + "a catalog names branch '" + branch + "' as it names the main branch"; + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, branch) + .Finish(), + message); + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch(branch) + .Finish(), + message); + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1$branch_" + branch)) + .Finish(), + message); + } + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch("dev$options") + .Finish(), + "a branch a catalog addresses cannot contain '$'"); + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1", "dev")) + .AddOption(Options::BRANCH, "test") + .Finish(), + "a read is aimed at one branch, but both 'dev' and 'test' were named"); + ASSERT_TRUE(catalog->TableFileSystemRequests().empty()); + ASSERT_EQ(0u, catalog->LoadTableSchemaCalls()); + ASSERT_OK_AND_ASSIGN(auto path_ctx, + ReadContextBuilder("table_root_path").WithBranch("MAIN").Finish()); + ASSERT_EQ("MAIN", path_ctx->GetBranch()); + + ASSERT_OK_AND_ASSIGN(auto option_ctx, ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ("dev", option_ctx->GetBranch()); + ASSERT_EQ(Identifier("db1", "t1$branch_dev"), option_ctx->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(auto branch_ctx, ReadContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch("dev") + .Finish()); + ASSERT_EQ(Identifier("db1", "t1$branch_dev"), branch_ctx->GetIdentifier().value()); + ASSERT_EQ(std::vector<Identifier>( + {Identifier("db1", "t1$branch_dev"), Identifier("db1", "t1$branch_dev")}), + catalog->TableFileSystemRequests()); +} + +TEST(ReadContextTest, TestSystemTablePathResolvesBranchBeforeCatalogRequest) { + auto catalog = std::make_shared<MockVersionManagedCatalog>(); + catalog->SetTableFileSystem(std::make_shared<MockFileSystem>()); + const std::string path = "table_root_path$branch_dev$files"; + for (const std::string table : {"t1", "t1$files"}) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> context, + ReadContextBuilder(path).WithCatalog(catalog, Identifier("db1", table)).Finish()); + ASSERT_EQ("dev", context->GetBranch()); + ASSERT_EQ(Identifier("db1", table == "t1" ? "t1$branch_dev" : "t1$branch_dev$files"), + context->GetIdentifier().value()); + ASSERT_EQ(context->GetIdentifier().value(), catalog->TableFileSystemRequests().back()); + } + const size_t requests = catalog->TableFileSystemRequests().size(); + ASSERT_NOK_WITH_MSG( + ReadContextBuilder(path).WithCatalog(catalog, Identifier("db1", "t1", "other")).Finish(), + "a read is aimed at one branch, but both 'other' and 'dev' were named"); + ASSERT_NOK_WITH_MSG(ReadContextBuilder(path) + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "main") + .Finish(), + "a read is aimed at one branch, but both 'main' and 'dev' were named"); + ASSERT_NOK_WITH_MSG(ReadContextBuilder(path) + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch("main") + .Finish(), + "a read is aimed at one branch, but both 'dev' and 'main' were named"); + ASSERT_NOK_WITH_MSG(ReadContextBuilder("table_root_path$branch_MAIN$files") + .WithCatalog(catalog, Identifier("db1", "t1")) + .Finish(), + "a catalog names branch 'MAIN' as it names the main branch"); + ASSERT_EQ(requests, catalog->TableFileSystemRequests().size()); +} + TEST(ReadContextTest, TestWithFileSystemOverridesCatalog) { // A file system the caller gave is the one that is used, and the catalog is not asked for one. auto table_fs = std::make_shared<MockFileSystem>();
diff --git a/src/paimon/core/operation/scan_context.cpp b/src/paimon/core/operation/scan_context.cpp index 89cccaa..87fb873 100644 --- a/src/paimon/core/operation/scan_context.cpp +++ b/src/paimon/core/operation/scan_context.cpp
@@ -22,6 +22,10 @@ #include "paimon/catalog/catalog.h" #include "paimon/common/utils/path_util.h" +#include "paimon/core/catalog/catalog_utils.h" +#include "paimon/core/table/system/system_table.h" +#include "paimon/core/utils/branch_manager.h" +#include "paimon/defs.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" #include "paimon/status.h" @@ -41,6 +45,23 @@ const std::map<std::string, std::string>& options, const std::shared_ptr<Cache>& cache, const std::shared_ptr<FormatTable>& format_table, const std::shared_ptr<Catalog>& catalog, const std::optional<Identifier>& identifier) + : ScanContext(path, is_streaming_mode, limit, scan_filter, global_index_result, + realtime_context, memory_pool, executor, specific_file_system, table_schema, + options, cache, format_table, catalog, identifier, + /*read_field_names=*/std::nullopt) {} + +ScanContext::ScanContext( + const std::string& path, bool is_streaming_mode, std::optional<int32_t> limit, + const std::shared_ptr<ScanFilter>& scan_filter, + const std::shared_ptr<GlobalIndexResult>& global_index_result, + const std::shared_ptr<RealtimeContext>& realtime_context, + const std::shared_ptr<MemoryPool>& memory_pool, const std::shared_ptr<Executor>& executor, + const std::shared_ptr<FileSystem>& specific_file_system, + const std::optional<std::string>& table_schema, + const std::map<std::string, std::string>& options, const std::shared_ptr<Cache>& cache, + const std::shared_ptr<FormatTable>& format_table, const std::shared_ptr<Catalog>& catalog, + const std::optional<Identifier>& identifier, + const std::optional<std::vector<std::string>>& read_field_names) : path_(path), is_streaming_mode_(is_streaming_mode), limit_(limit), @@ -55,7 +76,8 @@ cache_(cache), format_table_(format_table), catalog_(catalog), - identifier_(identifier) {} + identifier_(identifier), + read_field_names_(read_field_names) {} ScanContext::~ScanContext() = default; @@ -79,6 +101,7 @@ cache_.reset(); catalog_.reset(); identifier_.reset(); + read_field_names_.reset(); } private: @@ -103,6 +126,7 @@ std::shared_ptr<Cache> cache_; std::shared_ptr<Catalog> catalog_; std::optional<Identifier> identifier_; + std::optional<std::vector<std::string>> read_field_names_; }; ScanContextBuilder::ScanContextBuilder(const std::string& path) @@ -194,6 +218,12 @@ return *this; } +ScanContextBuilder& ScanContextBuilder::SetReadFieldNames( + const std::vector<std::string>& read_field_names) { + impl_->read_field_names_ = read_field_names; + return *this; +} + ScanContextBuilder& ScanContextBuilder::SetTableSchema(const std::string& table_schema) { impl_->table_schema_ = table_schema; return *this; @@ -234,6 +264,28 @@ if (impl_->catalog_ == nullptr && impl_->identifier_) { return Status::Invalid("cannot scan through a null catalog"); } + // Resolve a branch named by a system table path too, before requesting its credentials. + std::optional<SystemTablePath> system_table_path; + if (impl_->format_table_ == nullptr) { + PAIMON_ASSIGN_OR_RAISE(system_table_path, SystemTableLoader::TryParsePath(impl_->path_)); + } + PAIMON_ASSIGN_OR_RAISE( + std::string branch, + BranchManager::ResolveBranch(impl_->identifier_, impl_->options_, + system_table_path ? system_table_path->branch : std::nullopt, + "scan")); + if (impl_->catalog_ != nullptr && !BranchManager::IsMainBranch(branch)) { + PAIMON_RETURN_NOT_OK(BranchManager::CheckCatalogAddressableBranch(branch)); + } + if (branch != BranchManager::DEFAULT_MAIN_BRANCH) { + impl_->options_[Options::BRANCH] = branch; + } + if (impl_->identifier_) { + PAIMON_ASSIGN_OR_RAISE( + Identifier branch_identifier, + CatalogUtils::ResolvedBranchIdentifier(impl_->identifier_.value(), branch)); + impl_->identifier_.emplace(std::move(branch_identifier)); + } if (impl_->catalog_ != nullptr && impl_->specific_file_system_ == nullptr) { // A catalog issuing per-table temporary credentials only hands them out through // GetTableFileSystem, so the scan uses the credentials of this table. @@ -248,7 +300,7 @@ impl_->bucket_filter_), impl_->global_index_result_, impl_->realtime_context_, impl_->memory_pool_, executor, impl_->specific_file_system_, impl_->table_schema_, impl_->options_, impl_->cache_, - impl_->format_table_, impl_->catalog_, impl_->identifier_); + impl_->format_table_, impl_->catalog_, impl_->identifier_, impl_->read_field_names_); impl_->Reset(); return ctx; }
diff --git a/src/paimon/core/operation/scan_context_test.cpp b/src/paimon/core/operation/scan_context_test.cpp index 4ce2395..dea85d0 100644 --- a/src/paimon/core/operation/scan_context_test.cpp +++ b/src/paimon/core/operation/scan_context_test.cpp
@@ -70,6 +70,101 @@ ASSERT_TRUE(catalog->TableFileSystemRequests().empty()); } +TEST(ScanContextTest, TestWithCatalogResolvesTheBranch) { + auto catalog = std::make_shared<MockVersionManagedCatalog>(); + catalog->SetTableFileSystem(std::make_shared<MockFileSystem>()); + + for (const std::string branch : {"MAIN", "Main"}) { + const std::string message = + "a catalog names branch '" + branch + "' as it names the main branch"; + ASSERT_NOK_WITH_MSG(ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, branch) + .Finish(), + message); + ASSERT_NOK_WITH_MSG(ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1$branch_" + branch)) + .Finish(), + message); + } + ASSERT_NOK_WITH_MSG(ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "dev$options") + .Finish(), + "a branch a catalog addresses cannot contain '$'"); + ASSERT_NOK_WITH_MSG(ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1", "dev")) + .AddOption(Options::BRANCH, "test") + .Finish(), + "a scan is aimed at one branch, but both 'dev' and 'test' were named"); + ASSERT_TRUE(catalog->TableFileSystemRequests().empty()); + ASSERT_EQ(0u, catalog->LoadTableSchemaCalls()); + ASSERT_OK_AND_ASSIGN( + auto path_ctx, + ScanContextBuilder("table_root_path").AddOption(Options::BRANCH, "MAIN").Finish()); + ASSERT_EQ("MAIN", path_ctx->GetOptions().at(Options::BRANCH)); + + ASSERT_OK_AND_ASSIGN(auto ctx, ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1", "dev")) + .Finish()); + ASSERT_EQ("dev", ctx->GetOptions().at(Options::BRANCH)); + ASSERT_OK_AND_ASSIGN(auto option_ctx, ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ("dev", option_ctx->GetOptions().at(Options::BRANCH)); + ASSERT_EQ(Identifier("db1", "t1$branch_dev"), option_ctx->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(auto system_ctx, ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1$files")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ(Identifier("db1", "t1$branch_dev$files"), system_ctx->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(auto sys_ctx, ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("sys", "tables")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ(Identifier("sys", "tables"), sys_ctx->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(auto main_ctx, ScanContextBuilder("table_root_path") + .WithCatalog(catalog, Identifier("db1", "t1")) + .Finish()); + ASSERT_EQ(0u, main_ctx->GetOptions().count(Options::BRANCH)); + ASSERT_EQ(Identifier("db1", "t1"), main_ctx->GetIdentifier().value()); + ASSERT_EQ(std::vector<Identifier>({Identifier("db1", "t1$branch_dev"), + Identifier("db1", "t1$branch_dev"), + Identifier("db1", "t1$branch_dev$files"), + Identifier("sys", "tables"), Identifier("db1", "t1")}), + catalog->TableFileSystemRequests()); +} + +TEST(ScanContextTest, TestSystemTablePathResolvesBranchBeforeCatalogRequest) { + auto catalog = std::make_shared<MockVersionManagedCatalog>(); + catalog->SetTableFileSystem(std::make_shared<MockFileSystem>()); + const std::string path = "table_root_path$branch_dev$files"; + for (const std::string table : {"t1", "t1$files"}) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> context, + ScanContextBuilder(path).WithCatalog(catalog, Identifier("db1", table)).Finish()); + ASSERT_EQ("dev", context->GetOptions().at(Options::BRANCH)); + ASSERT_EQ(Identifier("db1", table == "t1" ? "t1$branch_dev" : "t1$branch_dev$files"), + context->GetIdentifier().value()); + ASSERT_EQ(context->GetIdentifier().value(), catalog->TableFileSystemRequests().back()); + } + const size_t requests = catalog->TableFileSystemRequests().size(); + ASSERT_NOK_WITH_MSG( + ScanContextBuilder(path).WithCatalog(catalog, Identifier("db1", "t1", "other")).Finish(), + "a scan is aimed at one branch, but both 'other' and 'dev' were named"); + ASSERT_NOK_WITH_MSG(ScanContextBuilder(path) + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "main") + .Finish(), + "a scan is aimed at one branch, but both 'main' and 'dev' were named"); + ASSERT_NOK_WITH_MSG(ScanContextBuilder("table_root_path$branch_MAIN$files") + .WithCatalog(catalog, Identifier("db1", "t1")) + .Finish(), + "a catalog names branch 'MAIN' as it names the main branch"); + ASSERT_EQ(requests, catalog->TableFileSystemRequests().size()); +} + TEST(ScanContextTest, TestWithCatalogNullRejected) { ScanContextBuilder builder("table_root_path"); ASSERT_NOK_WITH_MSG(builder.WithCatalog(nullptr, Identifier("db1", "t1")).Finish(), @@ -91,6 +186,7 @@ ASSERT_TRUE(ctx->GetOptions().empty()); ASSERT_FALSE(ctx->GetGlobalIndexResult()); ASSERT_FALSE(ctx->GetSpecificFileSystem()); + ASSERT_FALSE(ctx->GetReadFieldNames()); } TEST(ScanContextTest, TestSetContent) { @@ -117,6 +213,7 @@ builder.SetTableSchema("table-schema-json"); auto manifest_cache = std::make_shared<LruCache>(1024); builder.WithCache(manifest_cache); + builder.SetReadFieldNames({"f1", "f2"}); ASSERT_OK_AND_ASSIGN(auto ctx, builder.Finish()); ASSERT_EQ(ctx->GetPath(), "table_root_path"); ASSERT_TRUE(ctx->IsStreamingMode()); @@ -134,6 +231,8 @@ ASSERT_EQ(expected_options, ctx->GetOptions()); ASSERT_EQ(fs, ctx->GetSpecificFileSystem()); ASSERT_TRUE(ctx->GetCache()); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{"f1", "f2"}), + ctx->GetReadFieldNames()); } TEST(ScanContextTest, TestSetOptionsOverridesAddedOptions) {
diff --git a/src/paimon/core/schema/schema_validation.cpp b/src/paimon/core/schema/schema_validation.cpp index ffed109..a310ec9 100644 --- a/src/paimon/core/schema/schema_validation.cpp +++ b/src/paimon/core/schema/schema_validation.cpp
@@ -335,6 +335,17 @@ return Status::Invalid("a format table does not support real-time options"); } + // The schema's own value too, which the options of a call cannot disable. + PAIMON_ASSIGN_OR_RAISE( + bool schema_query_auth, + OptionsUtils::GetValueFromMap<bool>(schema.Options(), Options::QUERY_AUTH_ENABLED, false)); + if (schema_query_auth || core_options.QueryAuthEnabled()) { + return Status::Invalid(fmt::format( + "{} is not supported on a format table: its read does not apply row filters or column " + "masks.", + Options::QUERY_AUTH_ENABLED)); + } + // A format table has no branches: its data is the files under its one location, which is // where a read or a write would go whatever this named. Refused wherever the option comes // from - the schema, a catalog that parsed `tbl$branch_dev`, or one call's options - rather @@ -657,6 +668,35 @@ return Status::OK(); } +Result<bool> SchemaValidation::QueryAuthReadEnabled( + const std::map<std::string, std::string>& schema_options, + const std::map<std::string, std::string>& dynamic_options) { + PAIMON_ASSIGN_OR_RAISE( + bool schema_enabled, + OptionsUtils::GetValueFromMap<bool>(schema_options, Options::QUERY_AUTH_ENABLED, false)); + PAIMON_ASSIGN_OR_RAISE(bool enabled, + OptionsUtils::GetValueFromMap<bool>( + dynamic_options, Options::QUERY_AUTH_ENABLED, schema_enabled)); + if (schema_enabled && !enabled) { + return Status::Invalid(fmt::format("Option '{}' is not allowed to be disabled dynamically", + Options::QUERY_AUTH_ENABLED)); + } + return enabled; +} + +Status SchemaValidation::RejectQueryAuth(const TableSchema& schema, + const std::map<std::string, std::string>& dynamic_options, + const std::string& what) { + PAIMON_ASSIGN_OR_RAISE(bool enabled, QueryAuthReadEnabled(schema.Options(), dynamic_options)); + if (!enabled) { + return Status::OK(); + } + return Status::NotImplemented( + fmt::format("{} is not supported on a table with '{}=true': it would not apply the row " + "filters and column masks the catalog defines for the table.", + what, Options::QUERY_AUTH_ENABLED)); +} + Status SchemaValidation::ValidateRealtime(const TableSchema& schema, const CoreOptions& options) { const std::vector<std::string>& deduplicate_key_fields = options.GetRealtimeDeduplicateKeyFields();
diff --git a/src/paimon/core/schema/schema_validation.h b/src/paimon/core/schema/schema_validation.h index 6c8af1b..0ad605c 100644 --- a/src/paimon/core/schema/schema_validation.h +++ b/src/paimon/core/schema/schema_validation.h
@@ -27,6 +27,7 @@ #include "paimon/core/core_options.h" #include "paimon/core/schema/table_schema.h" +#include "paimon/result.h" #include "paimon/status.h" #include "paimon/visibility.h" @@ -55,6 +56,20 @@ /// table schema options. static Status ValidateRealtime(const TableSchema& schema, const CoreOptions& options); + /// Returns whether a read given `dynamic_options` of a table whose schema has + /// `schema_options` enables `query-auth.enabled`. Fails when `dynamic_options` disable what + /// `schema_options` enable. + static Result<bool> QueryAuthReadEnabled( + const std::map<std::string, std::string>& schema_options, + const std::map<std::string, std::string>& dynamic_options); + + /// Refuses `what`, a way of reading a table that does not apply the catalog's row filters and + /// column masks, when the schema or `dynamic_options` enable `query-auth.enabled`. + /// `dynamic_options` cannot disable what the schema enables. + static Status RejectQueryAuth(const TableSchema& schema, + const std::map<std::string, std::string>& dynamic_options, + const std::string& what); + /// Checks the deliberately limited dedicated-vector write path, including runtime overrides. static Status ValidateVectorStore(const TableSchema& schema, const CoreOptions& options);
diff --git a/src/paimon/core/schema/schema_validation_test.cpp b/src/paimon/core/schema/schema_validation_test.cpp index 44aaa71..f88601d 100644 --- a/src/paimon/core/schema/schema_validation_test.cpp +++ b/src/paimon/core/schema/schema_validation_test.cpp
@@ -109,6 +109,50 @@ ASSERT_OK(SchemaValidation::ValidateTableSchema(*table_schema)); } +TEST(SchemaValidationTest, TestQueryAuth) { + auto schema = arrow::schema({arrow::field("f0", arrow::int32())}); + auto create = [&](const std::map<std::string, std::string>& options) { + return TableSchema::Create(/*schema_id=*/0, schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options); + }; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableSchema> plain, create({})); + ASSERT_OK_AND_ASSIGN(bool plain_enabled, + SchemaValidation::QueryAuthReadEnabled(plain->Options(), {})); + ASSERT_FALSE(plain_enabled); + ASSERT_OK_AND_ASSIGN(bool enabled_by_read, + SchemaValidation::QueryAuthReadEnabled( + plain->Options(), {{Options::QUERY_AUTH_ENABLED, "true"}})); + ASSERT_TRUE(enabled_by_read); + ASSERT_OK(SchemaValidation::RejectQueryAuth(*plain, {}, "Search")); + ASSERT_OK(SchemaValidation::RejectQueryAuth(*plain, {{Options::QUERY_AUTH_ENABLED, "false"}}, + "Search")); + Status refused_by_read = SchemaValidation::RejectQueryAuth( + *plain, {{Options::QUERY_AUTH_ENABLED, "true"}}, "Search"); + ASSERT_TRUE(refused_by_read.IsNotImplemented()) << refused_by_read.ToString(); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableSchema> query_auth, + create({{Options::QUERY_AUTH_ENABLED, "true"}})); + ASSERT_OK(SchemaValidation::ValidateNewTableSchema(*query_auth, /*file_system=*/nullptr)); + ASSERT_OK_AND_ASSIGN(bool schema_enabled, + SchemaValidation::QueryAuthReadEnabled(query_auth->Options(), {})); + ASSERT_TRUE(schema_enabled); + ASSERT_NOK_WITH_MSG(SchemaValidation::RejectQueryAuth(*query_auth, {}, "Search"), + "NotImplemented: Search is not supported on a table with " + "'query-auth.enabled=true'"); + Status disabled = SchemaValidation::RejectQueryAuth( + *query_auth, {{Options::QUERY_AUTH_ENABLED, "false"}}, "Search"); + ASSERT_TRUE(disabled.IsInvalid()) << disabled.ToString(); + ASSERT_NOK_WITH_MSG(disabled, "is not allowed to be disabled dynamically"); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableSchema> format_table, + create({{Options::TYPE, "format-table"}, {Options::QUERY_AUTH_ENABLED, "true"}})); + ASSERT_NOK_WITH_MSG( + SchemaValidation::ValidateNewTableSchema(*format_table, /*file_system=*/nullptr), + "Invalid: query-auth.enabled is not supported on a format table"); +} + TEST(SchemaValidationTest, TestRealtime) { const std::shared_ptr<arrow::Schema> schema = arrow::schema({arrow::field("id", arrow::int64(), /*nullable=*/false)});
diff --git a/src/paimon/core/table/format/format_table_test.cpp b/src/paimon/core/table/format/format_table_test.cpp index 40e4c3e..9fb2fc4 100644 --- a/src/paimon/core/table/format/format_table_test.cpp +++ b/src/paimon/core/table/format/format_table_test.cpp
@@ -796,6 +796,45 @@ "metastore.partitioned-table"); } +TEST(FormatTableTest, TestQueryAuthInTheSchemaCannotBeDisabledAtTheCall) { + std::unique_ptr<UniqueTestDirectory> dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> schema, + TableSchema::Create(/*schema_id=*/0, MakeSchema(), /*partition_keys=*/{}, + /*primary_keys=*/{}, + {{Options::TYPE, "format-table"}, + {Options::FILE_FORMAT, "parquet"}, + {Options::QUERY_AUTH_ENABLED, "true"}})); + ASSERT_OK_AND_ASSIGN(std::string schema_json, schema->ToJsonString()); + SchemaManager schema_manager(dir->GetFileSystem(), dir->Str()); + ASSERT_OK(dir->GetFileSystem()->Mkdirs(schema_manager.SchemaDirectory())); + ASSERT_OK(dir->GetFileSystem()->WriteFile( + PathUtil::JoinPath(schema_manager.SchemaDirectory(), "schema-0"), schema_json, + /*overwrite=*/false)); + const std::map<std::string, std::string> disabling = {{Options::QUERY_AUTH_ENABLED, "false"}}; + const std::string refused = "query-auth.enabled is not supported on a format table"; + + Status created = + FormatTable::Create(dir->GetFileSystem(), dir->Str(), Identifier("db", "tbl"), disabling) + .status(); + ASSERT_TRUE(created.IsInvalid()) << created.ToString(); + ASSERT_NOK_WITH_MSG(created, refused); + + ScanContextBuilder scan_builder(dir->Str()); + scan_builder.SetOptions(disabling); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, scan_builder.Finish()); + Status scanned = TableScan::Create(std::move(scan_context)).status(); + ASSERT_TRUE(scanned.IsInvalid()) << scanned.ToString(); + ASSERT_NOK_WITH_MSG(scanned, refused); + + ReadContextBuilder read_builder(dir->Str()); + read_builder.SetOptions(disabling); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, read_builder.Finish()); + Status read = TableRead::Create(std::move(read_context)).status(); + ASSERT_TRUE(read.IsInvalid()) << read.ToString(); + ASSERT_NOK_WITH_MSG(read, refused); +} + TEST(FormatTableTest, TestValueOnlyPartitionCannotBeNamedAfterThisTablesMetadata) { // Under the value-only layout a value becomes a directory name unchanged, so `dt` of // `schema` would be written over this table's own schema, and an overwrite would delete it.
diff --git a/src/paimon/core/table/source/query_auth_split.h b/src/paimon/core/table/source/query_auth_split.h new file mode 100644 index 0000000..65670f1 --- /dev/null +++ b/src/paimon/core/table/source/query_auth_split.h
@@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <cstdint> +#include <map> +#include <memory> +#include <optional> +#include <string> +#include <utility> + +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/table/source/split.h" + +namespace paimon { + +/// The table, branch and columns a `QueryAuthSplit` was authorized for. +struct QueryAuthScope { + /// The table location and branch the scan authorized. Matching field ids alone must not let a + /// read apply the rules to a different table or branch with the same schema and file names. + std::string table_path; + std::string branch; + /// The field id, in the planned schema, of each column the scan asked the catalog to + /// authorize; every column of that schema when the scan asked for all of them. A read may + /// only output or filter on these columns, and only while they keep these field ids. + std::map<std::string, int32_t> authorized_field_ids; + /// The field id, in the planned schema, of each column the rules reference: the row filter + /// operands and the targets and inputs of the masks. A read refuses rules referencing a column + /// this does not bind, or one the schema it reads exposes under another field id. + std::map<std::string, int32_t> rule_field_ids; + + bool operator==(const QueryAuthScope& other) const { + return table_path == other.table_path && branch == other.branch && + authorized_field_ids == other.authorized_field_ids && + rule_field_ids == other.rule_field_ids; + } +}; + +/// A split of a table with `query-auth.enabled=true`, carrying the rules the scan that planned it +/// got from the catalog, so the read applies them. Mirrors Java `QueryAuthSplit`. +/// +/// A split planned here carries its `QueryAuthScope` and serializes in a paimon-cpp envelope, +/// which Java rejects as a corrupted split. A split Java serialized has no scope, so the read +/// refuses it; written back, it takes the Java `QueryAuthSplit` format. +class QueryAuthSplit : public Split { + public: + /// The magic (ASCII "SPLIT_V1") and version Java `SplitSerializer` starts a split with, and + /// its type tag. + static constexpr int64_t kSplitSerializerMagic = 0x53504C49545F5631LL; + static constexpr int32_t kSplitSerializerVersion = 1; + static constexpr int32_t kSplitSerializerDataSplit = 1; + static constexpr int32_t kSplitSerializerIndexedSplit = 3; + static constexpr int32_t kSplitSerializerQueryAuthSplit = 5; + + /// The magic and version of the envelope a split carrying its scope is written in: the Java + /// `SplitSerializer` bytes of the split, then the scope. Both parts are required. The magic is + /// ASCII "PCQAUTH1". + static constexpr int64_t kScopedMagic = 0x5043514155544831LL; + static constexpr int32_t kScopedVersion = 1; + + /// @param scope What the rules were authorized for, or null when the split does not record it. + QueryAuthSplit(std::shared_ptr<Split> split, std::shared_ptr<TableQueryAuthResult> auth_result, + std::optional<QueryAuthScope> scope) + : split_(std::move(split)), + auth_result_(std::move(auth_result)), + scope_(std::move(scope)) {} + + const std::shared_ptr<Split>& GetSplit() const { + return split_; + } + + const std::shared_ptr<TableQueryAuthResult>& GetAuthResult() const { + return auth_result_; + } + + const std::optional<QueryAuthScope>& GetScope() const { + return scope_; + } + + private: + std::shared_ptr<Split> split_; + std::shared_ptr<TableQueryAuthResult> auth_result_; + std::optional<QueryAuthScope> scope_; +}; + +} // namespace paimon
diff --git a/src/paimon/core/table/source/query_auth_split_test.cpp b/src/paimon/core/table/source/query_auth_split_test.cpp new file mode 100644 index 0000000..777eb1d --- /dev/null +++ b/src/paimon/core/table/source/query_auth_split_test.cpp
@@ -0,0 +1,359 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/table/source/query_auth_split.h" + +#include <cstdint> +#include <map> +#include <memory> +#include <optional> +#include <set> +#include <string> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "gtest/gtest.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/core/global_index/indexed_split_impl.h" +#include "paimon/core/io/data_file_meta.h" +#include "paimon/core/table/source/data_split_impl.h" +#include "paimon/data/timestamp.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/result.h" +#include "paimon/status.h" +#include "paimon/table/source/split.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/range.h" + +namespace paimon::test { + +namespace { + +Result<std::string> ReadFixture(const std::string& name) { + LocalFileSystem file_system; + std::string bytes; + PAIMON_RETURN_NOT_OK(file_system.ReadFile(GetDataDir() + "/query_auth/" + name, &bytes)); + return bytes; +} + +Result<std::shared_ptr<DataSplitImpl>> FixtureDataSplit() { + auto file = std::make_shared<DataFileMeta>( + "file1.orc", 100l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(), + SimpleStats::EmptyStats(), SimpleStats::EmptyStats(), 50l, 249l, 0, 0, + std::vector<std::optional<std::string>>(), Timestamp(1765535214349l, 0), 0, nullptr, + FileSource::Append(), std::nullopt, std::nullopt, 50l, std::nullopt, + /*column_max_sequence_numbers=*/std::nullopt); + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, "data/test_table/bucket-0", + std::vector<std::shared_ptr<DataFileMeta>>({file})); + return builder.WithSnapshot(1).IsStreaming(false).RawConvertible(true).Build(); +} + +} // namespace + +TEST(QueryAuthSplitTest, TestDeserializeSplitsGeneratedByJava) { + auto pool = GetDefaultPool(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> expected_data_split, FixtureDataSplit()); + const std::string id_ref = R"({"index": 0, "name": "id", "type": "INT NOT NULL"})"; + const std::string name_ref = R"({"index": 1, "name": "name", "type": "STRING"})"; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> expected_rules, + TableQueryAuthResult::Create( + {fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "FIELD_REF", "fieldRef": )" + R"({}}}, "function": "GREATER_THAN", "literals": [1]}})", + id_ref), + fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "FIELD_REF", "fieldRef": )" + R"({}}}, "function": "NOT_EQUAL", "literals": ["测试"]}})", + name_ref)}, + {})); + auto batch = checked_pointer_cast<arrow::StructArray>( + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_({arrow::field("id", arrow::int32(), /*nullable=*/false), + arrow::field("name", arrow::utf8()), + arrow::field("phone", arrow::utf8()), + arrow::field("salary", arrow::int64())}), + R"([[2, "bob", "13900002222", 200]])") + .ValueOrDie()); + + for (const std::string name : + {"query_auth_split-data_split", "query_auth_split-indexed_split"}) { + SCOPED_TRACE(name); + ASSERT_OK_AND_ASSIGN(std::string bytes, ReadFixture(name)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> split, + Split::Deserialize(bytes.data(), bytes.size(), pool)); + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split); + ASSERT_NE(nullptr, auth_split); + ASSERT_FALSE(auth_split->GetScope()); + if (auto indexed = std::dynamic_pointer_cast<IndexedSplitImpl>(auth_split->GetSplit())) { + ASSERT_TRUE(indexed->TEST_Equal(IndexedSplitImpl( + expected_data_split, std::vector<Range>{Range(55, 56)}, std::vector<float>{0.5f}))); + } else { + auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(auth_split->GetSplit()); + ASSERT_NE(nullptr, data_split); + ASSERT_TRUE(data_split->TEST_Equal(*expected_data_split)); + } + const TableQueryAuthResult& rules = *auth_split->GetAuthResult(); + ASSERT_TRUE(rules.SameFilter(*expected_rules)); + ASSERT_EQ((std::set<std::string>{"name", "phone", "salary"}), rules.MaskedFields()); + for (const auto& [column, expected] : std::vector<std::pair<std::string, std::string>>{ + {"name", R"(["名字"])"}, {"phone", R"(["139"])"}}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> masked, + rules.Masks().at(column)->Evaluate(*batch, arrow::utf8(), + arrow::default_memory_pool())); + ASSERT_TRUE(masked->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), expected).ValueOrDie())) + << column << ": " << masked->ToString(); + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> salary, + rules.Masks().at("salary")->Evaluate(*batch, arrow::int64(), + arrow::default_memory_pool())); + ASSERT_EQ(1, salary->null_count()); + ASSERT_OK_AND_ASSIGN(std::string written, Split::Serialize(split, pool)); + ASSERT_EQ(bytes, written); + const std::string with_tail = bytes + std::string(8, '\0'); + ASSERT_NOK_WITH_MSG(Split::Deserialize(with_tail.data(), with_tail.size(), pool), + "remaining 8 bytes after deserializing"); + } + + ASSERT_OK_AND_ASSIGN(std::string empty_rules, ReadFixture("query_auth_split-empty_rules")); + ASSERT_OK_AND_ASSIGN(std::string null_rules, ReadFixture("query_auth_split-null_rules")); + for (const std::string* bytes : {&empty_rules, &null_rules}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> split, + Split::Deserialize(bytes->data(), bytes->size(), pool)); + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split); + ASSERT_NE(nullptr, auth_split); + ASSERT_FALSE(auth_split->GetAuthResult()->HasRules()); + auto data_split = std::dynamic_pointer_cast<DataSplitImpl>(auth_split->GetSplit()); + ASSERT_NE(nullptr, data_split); + ASSERT_TRUE(data_split->TEST_Equal(*expected_data_split)); + ASSERT_OK_AND_ASSIGN(std::string written, Split::Serialize(split, pool)); + ASSERT_EQ(empty_rules, written); + } + ASSERT_NE(empty_rules, null_rules); + + const size_t empty_auth_size = 3 + 2 * sizeof(int32_t); + for (int8_t flag : {static_cast<int8_t>(2), static_cast<int8_t>(-1)}) { + std::string noncanonical_flags = empty_rules; + for (size_t offset : {empty_auth_size, empty_auth_size - 1, sizeof(int32_t) + 1}) { + noncanonical_flags[noncanonical_flags.size() - offset] = static_cast<char>(flag); + } + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<Split> split, + Split::Deserialize(noncanonical_flags.data(), noncanonical_flags.size(), pool)); + ASSERT_OK_AND_ASSIGN(std::string written, Split::Serialize(split, pool)); + ASSERT_EQ(empty_rules, written); + } + + ASSERT_OK_AND_ASSIGN(std::string null_auth_result, + ReadFixture("query_auth_split-null_auth_result")); + ASSERT_NOK_WITH_MSG(Split::Deserialize(null_auth_result.data(), null_auth_result.size(), pool), + "invalid QueryAuthSplit, it carries no rules"); +} + +TEST(QueryAuthSplitTest, TestReadsNestedIndexedSplitAsJava) { + ASSERT_OK_AND_ASSIGN(std::string fixture, ReadFixture("query_auth_split-indexed_split")); + auto pool = GetDefaultPool(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> split, + Split::Deserialize(fixture.data(), fixture.size(), pool)); + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split); + ASSERT_NE(nullptr, auth_split); + auto indexed = std::dynamic_pointer_cast<IndexedSplitImpl>(auth_split->GetSplit()); + ASSERT_NE(nullptr, indexed); + ASSERT_EQ(1u, indexed->RowRanges().size()); + ASSERT_EQ(1u, indexed->Scores().size()); + ASSERT_OK_AND_ASSIGN(std::string data_bytes, Split::Serialize(indexed->GetDataSplit(), pool)); + for (const std::optional<QueryAuthScope>& scope : + {std::optional<QueryAuthScope>(), std::optional<QueryAuthScope>(QueryAuthScope())}) { + ASSERT_OK_AND_ASSIGN(std::string bytes, + Split::Serialize(std::make_shared<QueryAuthSplit>( + indexed, auth_split->GetAuthResult(), scope), + pool)); + const size_t range_count_offset = (scope ? 12 : 0) + 2 * 16 + 12 + data_bytes.size(); + const size_t score_count_offset = range_count_offset + 4 + 2 * sizeof(int64_t) + 1; + for (size_t offset : {range_count_offset, score_count_offset}) { + for (const std::string& count : + {std::string(4, '\xFF'), std::string("\x00\x00\x04\x00", 4)}) { + SCOPED_TRACE(offset); + std::string invalid = bytes; + invalid.replace(offset, count.size(), count); + ASSERT_NOK_WITH_MSG(Split::Deserialize(invalid.data(), invalid.size(), pool), + "invalid IndexedSplit"); + } + } + } + + auto original_data = std::dynamic_pointer_cast<DataSplitImpl>(indexed->GetDataSplit()); + ASSERT_NE(nullptr, original_data); + DataSplitImpl::Builder builder( + original_data->Partition(), original_data->Bucket(), original_data->BucketPath(), + std::vector<std::shared_ptr<DataFileMeta>>(original_data->DataFiles())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> data, + builder.WithSnapshot(original_data->SnapshotId()) + .WithTotalBuckets(8) + .IsStreaming(true) + .RawConvertible(true) + .Build()); + auto with_total_buckets = + std::make_shared<IndexedSplitImpl>(data, indexed->RowRanges(), indexed->Scores()); + ASSERT_OK_AND_ASSIGN( + std::string bytes, + Split::Serialize(std::make_shared<QueryAuthSplit>( + with_total_buckets, auth_split->GetAuthResult(), std::nullopt), + pool)); + ASSERT_OK_AND_ASSIGN(data_bytes, Split::Serialize(data, pool)); + const size_t data_start = 2 * 16 + 12; + const size_t data_end = data_start + data_bytes.size(); + const size_t has_total_buckets_offset = + data_start + 20 + 2 * sizeof(int32_t) + original_data->Partition().GetSizeInBytes() + + sizeof(int32_t) + sizeof(uint16_t) + original_data->BucketPath().size(); + const size_t has_scores_offset = data_end + 4 + 2 * sizeof(int64_t); + for (char flag : {'\x02', '\xFF'}) { + std::string noncanonical = bytes; + noncanonical[has_total_buckets_offset] = flag; + noncanonical[data_end - 2] = flag; + noncanonical[data_end - 1] = flag; + noncanonical[has_scores_offset] = flag; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> parsed, + Split::Deserialize(noncanonical.data(), noncanonical.size(), pool)); + ASSERT_OK_AND_ASSIGN(std::string written, Split::Serialize(parsed, pool)); + ASSERT_EQ(bytes, written); + } +} + +TEST(QueryAuthSplitTest, TestRoundTripsAndRefusesMalformedBytes) { + auto pool = GetDefaultPool(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<DataSplitImpl> data_split, FixtureDataSplit()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> rules, + TableQueryAuthResult::Create( + {R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 3,)" + R"( "name": "salary", "type": "BIGINT"}}, "function": "GREATER_THAN",)" + R"( "literals": [150]})"}, + {{"name", R"({"name": "CONCAT", "inputs": ["名字"]})"}})); + QueryAuthScope planned_scope; + planned_scope.table_path = "file:///warehouse/db/table"; + planned_scope.branch = "dev"; + planned_scope.authorized_field_ids = {{"id", 0}, {"name", 1}, {"phone", 2}, {"salary", 3}}; + planned_scope.rule_field_ids = {{"name", 1}, {"salary", 3}}; + auto planned = std::make_shared<QueryAuthSplit>(data_split, rules, planned_scope); + + auto indexed_split = std::make_shared<IndexedSplitImpl>( + data_split, std::vector<Range>{Range(0, 1)}, std::vector<float>{0.5f}); + for (const std::optional<QueryAuthScope>& scope : + {planned->GetScope(), std::optional<QueryAuthScope>()}) { + ASSERT_OK_AND_ASSIGN(std::string bytes, + Split::Serialize(std::make_shared<QueryAuthSplit>( + indexed_split, planned->GetAuthResult(), scope), + pool)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> copy, + Split::Deserialize(bytes.data(), bytes.size(), pool)); + auto auth_copy = std::dynamic_pointer_cast<QueryAuthSplit>(copy); + ASSERT_NE(nullptr, auth_copy); + auto indexed_copy = std::dynamic_pointer_cast<IndexedSplitImpl>(auth_copy->GetSplit()); + ASSERT_NE(nullptr, indexed_copy); + ASSERT_TRUE(indexed_copy->TEST_Equal(*indexed_split)); + ASSERT_TRUE(auth_copy->GetAuthResult()->SameFilter(*planned->GetAuthResult())); + ASSERT_EQ(planned->GetAuthResult()->ColumnMasking(), + auth_copy->GetAuthResult()->ColumnMasking()); + ASSERT_EQ(scope, auth_copy->GetScope()); + } + + ASSERT_NOK_WITH_MSG(Split::Serialize(std::make_shared<QueryAuthSplit>( + planned, planned->GetAuthResult(), std::nullopt), + pool), + "a QueryAuthSplit can only wrap a DataSplit or an IndexedSplit"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> no_rules_result, + TableQueryAuthResult::Create({}, {})); + ASSERT_OK_AND_ASSIGN( + std::string no_rules, + Split::Serialize( + std::make_shared<QueryAuthSplit>(data_split, no_rules_result, QueryAuthScope()), pool)); + ASSERT_OK(Split::Deserialize(no_rules.data(), no_rules.size(), pool).status()); + const size_t envelope_header_size = 12; + const size_t empty_scope_size = 4 * sizeof(int32_t); + const std::string java_format_bytes = no_rules.substr( + envelope_header_size, no_rules.size() - envelope_header_size - empty_scope_size); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<Split> java_split, + Split::Deserialize(java_format_bytes.data(), java_format_bytes.size(), pool)); + ASSERT_FALSE(std::dynamic_pointer_cast<QueryAuthSplit>(java_split)->GetScope()); + ASSERT_OK_AND_ASSIGN(std::string java_copy, Split::Serialize(java_split, pool)); + ASSERT_EQ(java_format_bytes, java_copy); + const size_t split_serializer_header_size = 16; + std::string nested; + for (int32_t i = 0; i < 100000; ++i) { + nested += java_format_bytes.substr(0, split_serializer_header_size); + } + nested += java_format_bytes.substr(split_serializer_header_size); + ASSERT_NOK_WITH_MSG(Split::Deserialize(nested.data(), nested.size(), pool), + "invalid QueryAuthSplit, it wraps another QueryAuthSplit"); + + ASSERT_NOK(Split::Deserialize(no_rules.data(), no_rules.size() - empty_scope_size, pool)); + ASSERT_NOK(Split::Deserialize(no_rules.data(), no_rules.size() - 1, pool)); + std::string negative_count = no_rules; + negative_count.replace(negative_count.size() - empty_scope_size, 4, std::string(4, '\xFF')); + ASSERT_NOK_WITH_MSG(Split::Deserialize(negative_count.data(), negative_count.size(), pool), + "invalid QueryAuthSplit, -1 entries do not fit the split"); + std::string oversized_count = no_rules; + oversized_count.replace(oversized_count.size() - empty_scope_size, 4, + std::string("\x00\x00\x01\x00", 4)); + ASSERT_NOK_WITH_MSG(Split::Deserialize(oversized_count.data(), oversized_count.size(), pool), + "invalid QueryAuthSplit, 256 entries do not fit the split"); + std::string unknown_version = no_rules; + unknown_version.replace(8, 4, std::string("\x00\x00\x00\x02", 4)); + ASSERT_NOK_WITH_MSG(Split::Deserialize(unknown_version.data(), unknown_version.size(), pool), + "Unsupported QueryAuthSplit version: 2"); + const std::string entry( + "\x00\x00\x00\x01" + "a\x00\x00\x00\x00", + 9); + const std::string duplicate_column = no_rules.substr(0, no_rules.size() - empty_scope_size) + + std::string("\x00\x00\x00\x02", 4) + entry + entry + + std::string(3 * sizeof(int32_t), '\0'); + ASSERT_NOK_WITH_MSG(Split::Deserialize(duplicate_column.data(), duplicate_column.size(), pool), + "invalid QueryAuthSplit, its scope names a column twice"); + auto auth_string = [](const std::string& value) { + const auto size = static_cast<uint32_t>(value.size()); + std::string bytes; + for (int32_t shift = 24; shift >= 0; shift -= 8) { + bytes.push_back(static_cast<char>(size >> shift)); + } + return bytes + value; + }; + const std::string mask_entry = auth_string("name") + auth_string(R"({"name": "NULL"})"); + const std::string without_mask_count = + java_format_bytes.substr(0, java_format_bytes.size() - sizeof(int32_t)); + const std::string duplicate_mask = + without_mask_count + std::string("\x00\x00\x00\x02", 4) + mask_entry + mask_entry; + ASSERT_NOK_WITH_MSG(Split::Deserialize(duplicate_mask.data(), duplicate_mask.size(), pool), + "invalid QueryAuthSplit, its rules mask a column twice"); + const std::string one_mask = + without_mask_count + std::string("\x00\x00\x00\x01", 4) + mask_entry; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> masked_split, + Split::Deserialize(one_mask.data(), one_mask.size(), pool)); + ASSERT_EQ( + (std::map<std::string, std::string>{{"name", R"({"name": "NULL"})"}}), + std::dynamic_pointer_cast<QueryAuthSplit>(masked_split)->GetAuthResult()->ColumnMasking()); +} + +} // namespace paimon::test
diff --git a/src/paimon/core/table/source/query_auth_table_read.cpp b/src/paimon/core/table/source/query_auth_table_read.cpp new file mode 100644 index 0000000..b004dda --- /dev/null +++ b/src/paimon/core/table/source/query_auth_table_read.cpp
@@ -0,0 +1,524 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/table/source/query_auth_table_read.h" + +#include <cstdint> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "arrow/compute/cast.h" +#include "fmt/format.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/reader/predicate_batch_reader.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/object_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/casting/casting_utils.h" +#include "paimon/core/catalog/catalog_utils.h" +#include "paimon/core/table/source/query_auth_split.h" +#include "paimon/core/utils/branch_manager.h" +#include "paimon/defs.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +/// The column as a mask or a predicate reads it: dictionary-decoded, and cast to a string when it +/// is a large string. +Result<std::shared_ptr<arrow::Array>> DecodeDictionaryAndLargeString( + const std::shared_ptr<arrow::Array>& column, arrow::MemoryPool* pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> decoded, + CastingUtils::DecodeDictionary(column, pool)); + if (decoded->type_id() == arrow::Type::LARGE_STRING) { + return CastingUtils::Cast(decoded, arrow::utf8(), arrow::compute::CastOptions::Safe(), + pool); + } + return decoded; +} + +/// The type `DecodeDictionaryAndLargeString()` gives a column of `type`. +std::shared_ptr<arrow::DataType> DecodedType(const std::shared_ptr<arrow::DataType>& type) { + std::shared_ptr<arrow::DataType> value_type = type; + if (value_type->id() == arrow::Type::DICTIONARY) { + value_type = checked_cast<const arrow::DictionaryType&>(*value_type).value_type(); + } + if (value_type->id() == arrow::Type::LARGE_STRING) { + return arrow::utf8(); + } + return value_type; +} + +/// Hands out only the rows the row filter keeps. `PredicateBatchReader::NextBatchWithBitmap()` +/// only clears the bits of the rows it drops and leaves their values in the Arrow array, so every +/// batch is compacted through `NextBatch()`. +class RowFilterBoundaryReader : public BatchReader { + public: + explicit RowFilterBoundaryReader(std::unique_ptr<PredicateBatchReader> reader) + : reader_(std::move(reader)) {} + + Result<ReadBatch> NextBatch() override { + return reader_->NextBatch(); + } + + Result<ReadBatchWithBitmap> NextBatchWithBitmap() override { + return BatchReader::NextBatchWithBitmap(); + } + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + private: + std::unique_ptr<PredicateBatchReader> reader_; +}; + +} // namespace + +QueryAuthBatchReader::QueryAuthBatchReader( + std::unique_ptr<BatchReader> reader, + std::map<std::string, std::shared_ptr<MaskTransform>> masks, + std::set<std::string> dropped_fields, std::set<std::string> decoded_fields, + std::shared_ptr<arrow::MemoryPool> arrow_pool, std::shared_ptr<arrow::DataType> logical_type) + : reader_(std::move(reader)), + masks_(std::move(masks)), + dropped_fields_(std::move(dropped_fields)), + decoded_fields_(std::move(decoded_fields)), + arrow_pool_(std::move(arrow_pool)), + logical_type_(std::move(logical_type)) { + decode_inputs_ = decoded_fields_; + for (const auto& [target, mask] : masks_) { + for (const std::string& input : mask->InputFields()) { + decode_inputs_.insert(input); + } + } +} + +Result<std::shared_ptr<arrow::StructArray>> QueryAuthBatchReader::DecodeInputs( + const std::shared_ptr<arrow::StructArray>& batch) const { + arrow::FieldVector fields = batch->struct_type()->fields(); + arrow::ArrayVector children = batch->fields(); + bool decoded_any = false; + for (size_t i = 0; i < fields.size(); ++i) { + if (decode_inputs_.count(fields[i]->name()) == 0) { + continue; + } + Result<std::shared_ptr<arrow::Array>> result = + DecodeDictionaryAndLargeString(children[i], arrow_pool_.get()); + if (!result.ok()) { + return result.status().WithMessage( + fmt::format("Query auth read cannot decode column '{}': {}", fields[i]->name(), + result.status().message())); + } + std::shared_ptr<arrow::Array> decoded = std::move(result).value(); + if (decoded != children[i]) { + fields[i] = fields[i]->WithType(decoded->type()); + children[i] = std::move(decoded); + decoded_any = true; + } + } + if (!decoded_any) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> decoded_batch, + arrow::StructArray::Make(children, fields)); + return decoded_batch; +} + +Result<BatchReader::ReadBatch> QueryAuthBatchReader::NextBatch() { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(c_array.get(), c_schema.get())); + if (array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("a query-auth read requires struct batches"); + } + auto struct_array = checked_pointer_cast<arrow::StructArray>(array); + if (struct_array->null_count() != 0) { + return Status::Invalid("a query-auth read does not support null rows in a batch"); + } + if (logical_type_ != nullptr) { + Result<std::shared_ptr<arrow::StructArray>> result = + ArrowUtils::FlattenUnresolvableDictionaries( + struct_array, logical_type_, arrow_pool_.get(), + /*preserve_layout_recoverable_dictionaries=*/false); + if (!result.ok()) { + return result.status().WithMessage( + fmt::format("Query auth read cannot decode output dictionaries: {}", + result.status().message())); + } + struct_array = std::move(result).value(); + } + const auto& struct_type = *struct_array->struct_type(); + std::shared_ptr<arrow::StructArray> decoded = struct_array; + if (!decode_inputs_.empty()) { + PAIMON_ASSIGN_OR_RAISE(decoded, DecodeInputs(struct_array)); + } + arrow::FieldVector fields; + arrow::ArrayVector children; + for (int32_t i = 0; i < struct_type.num_fields(); ++i) { + const std::shared_ptr<arrow::Field>& field = struct_type.field(i); + if (dropped_fields_.count(field->name()) > 0) { + continue; + } + auto mask = masks_.find(field->name()); + if (mask == masks_.end()) { + const arrow::StructArray& source = + decoded_fields_.count(field->name()) > 0 ? *decoded : *struct_array; + fields.push_back(source.struct_type()->field(i)); + children.push_back(source.field(i)); + continue; + } + const std::shared_ptr<arrow::DataType> target_type = DecodedType(field->type()); + Result<std::shared_ptr<arrow::Array>> result = + mask->second->Evaluate(*decoded, target_type, arrow_pool_.get()); + if (!result.ok()) { + return result.status().WithMessage( + fmt::format("Query auth mask {} of column '{}' failed: {}", mask->second->Name(), + field->name(), result.status().message())); + } + std::shared_ptr<arrow::Array> masked = std::move(result).value(); + if (masked->length() != struct_array->length()) { + return Status::Invalid(fmt::format("the mask of column '{}' produced {} rows for {}", + field->name(), masked->length(), + struct_array->length())); + } + if (!masked->type()->Equals(*target_type)) { + return Status::Invalid(fmt::format("the mask of column '{}' produced {} for {}", + field->name(), masked->type()->ToString(), + target_type->ToString())); + } + // A mask may null a column that holds no nulls. + fields.push_back(field->WithType(target_type)->WithNullable(true)); + children.push_back(std::move(masked)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> output, + arrow::StructArray::Make(children, fields)); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*output, c_array.get(), c_schema.get())); + // The masked buffers come from `arrow_pool_`, which must outlive the batch, not this reader. + PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); + return batch; +} + +Result<std::unique_ptr<TableRead>> QueryAuthTableRead::Create( + const std::shared_ptr<ReadContext>& context, const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, InnerReadFactory create_inner) { + if (context->HasVectorOrFullTextSearch() || context->EnablePreciseFileIndexEvaluate()) { + return Status::NotImplemented( + "Search is not supported on a query-auth table: the index ranks raw values, which a " + "column mask invalidates."); + } + if (context->HasReadSchema()) { + return Status::NotImplemented( + "An explicit read schema is not supported on a query-auth table: the rules apply to " + "whole columns; use SetReadFieldNames() or SetReadFieldIds() instead."); + } + if (context->GetRealtimeContext() != nullptr) { + return Status::NotImplemented("Real-time reads are not supported on a query-auth table."); + } + std::optional<std::vector<std::string>> output_fields; + if (!context->GetReadFieldIds().empty()) { + std::vector<std::string> names; + for (int32_t id : context->GetReadFieldIds()) { + Result<DataField> field = table_schema->GetField(id); + if (!field.ok()) { + return Status::NotImplemented(fmt::format( + "Field {} is not a column of the table schema, which is all a query-auth read " + "can authorize.", + id)); + } + names.push_back(field.value().Name()); + } + output_fields = std::move(names); + } else if (!context->GetReadFieldNames().empty()) { + output_fields = context->GetReadFieldNames(); + } + std::set<std::string> predicate_fields; + if (context->GetPredicate() != nullptr) { + PAIMON_RETURN_NOT_OK( + PredicateUtils::GetAllNames(context->GetPredicate(), &predicate_fields)); + } + std::vector<std::string> requested = output_fields.value_or(std::vector<std::string>()); + requested.insert(requested.end(), predicate_fields.begin(), predicate_fields.end()); + for (const std::string& name : requested) { + if (!table_schema->GetField(name).ok()) { + return Status::NotImplemented(fmt::format( + "Column '{}' is not a column of the table schema, which is all a query-auth read " + "can authorize.", + name)); + } + } + std::unique_ptr<QueryAuthTableRead> read(new QueryAuthTableRead( + context, table_schema, arrow_pool, std::move(create_inner), std::move(output_fields))); + // Built now, so a read the table cannot serve fails here rather than at the first split. + PAIMON_RETURN_NOT_OK(read->GetOrCreateInnerRead({}).status()); + return std::unique_ptr<TableRead>(std::move(read)); +} + +QueryAuthTableRead::QueryAuthTableRead(const std::shared_ptr<ReadContext>& context, + const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, + InnerReadFactory create_inner, + std::optional<std::vector<std::string>> output_fields) + : context_(context), + table_schema_(table_schema), + arrow_pool_(arrow_pool), + create_inner_(std::move(create_inner)), + output_fields_(std::move(output_fields)) {} + +Result<std::shared_ptr<Predicate>> QueryAuthTableRead::PostMaskFilter( + const TableQueryAuthResult& rules) const { + const std::shared_ptr<Predicate>& predicate = context_->GetPredicate(); + if (predicate == nullptr || context_->EnablePredicateFilter()) { + return predicate; + } + // As Java `ReadTransform`: the scan did not prune by the conjuncts on masked columns, so the + // read evaluates them even when the caller evaluates the rest of its predicate itself. + const std::set<std::string> masked_fields = rules.MaskedFields(); + std::vector<std::shared_ptr<Predicate>> masked_conjuncts; + for (const std::shared_ptr<Predicate>& conjunct : PredicateUtils::SplitAnd(predicate)) { + PAIMON_ASSIGN_OR_RAISE(bool masked, + PredicateUtils::ContainAnyField(conjunct, masked_fields)); + if (masked) { + masked_conjuncts.push_back(conjunct); + } + } + if (masked_conjuncts.empty()) { + return std::shared_ptr<Predicate>(); + } + return PredicateBuilder::And(masked_conjuncts); +} + +Result<TableRead*> QueryAuthTableRead::GetOrCreateInnerRead( + const std::vector<std::string>& extra_fields) { + auto iter = inner_reads_.find(extra_fields); + if (iter != inner_reads_.end()) { + return iter->second.get(); + } + std::vector<std::string> read_field_names = context_->GetReadFieldNames(); + std::vector<int32_t> read_field_ids = context_->GetReadFieldIds(); + if (!read_field_ids.empty()) { + for (const std::string& name : extra_fields) { + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema_->GetField(name)); + read_field_ids.push_back(field.Id()); + } + } else if (!read_field_names.empty()) { + read_field_names.insert(read_field_names.end(), extra_fields.begin(), extra_fields.end()); + } + // Without the predicate: file statistics and indexes hold raw values. + auto inner_context = std::make_unique<ReadContext>( + context_->GetPath(), context_->GetBranch(), read_field_names, read_field_ids, + /*predicate=*/nullptr, /*vector_search=*/nullptr, /*full_text_search=*/nullptr, + /*enable_precise_file_index_evaluate=*/false, /*enable_predicate_filter=*/false, + context_->EnablePrefetch(), context_->EnableLateMaterializing(), + context_->GetPrefetchBatchCount(), context_->GetPrefetchMaxParallelNum(), + context_->EnableMultiThreadRowToBatch(), context_->GetRowToBatchThreadNumber(), + context_->GetSpecificTableSchema(), context_->GetMemoryPool(), context_->GetExecutor(), + context_->GetSpecificFileSystem(), context_->GetFileSystemSchemeToIdentifierMap(), + /*realtime_context=*/nullptr, context_->GetOptions(), context_->ReadAheadCacheEnabled(), + context_->GetCacheConfig(), context_->GetCache(), context_->GetFormatTable(), + context_->GetWarmupLevel(), context_->GetCatalog(), context_->GetIdentifier()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> read, + create_inner_(std::move(inner_context))); + TableRead* raw_read = read.get(); + inner_reads_.emplace(extra_fields, std::move(read)); + return raw_read; +} + +Status QueryAuthTableRead::ValidateSplit(const QueryAuthSplit& split, + const std::set<std::string>& filter_fields) const { + if (!split.GetScope()) { + return Status::Invalid( + "The split does not record which columns its query auth rules were authorized for, " + "as a split Java planned does not; refusing to read it. Plan the scan with " + "paimon-cpp."); + } + const QueryAuthScope& scope = split.GetScope().value(); + if (scope.table_path.empty() || scope.branch.empty()) { + return Status::Invalid( + "The split does not record the table path and branch its query auth rules were " + "authorized for. Plan the scan again."); + } + PAIMON_ASSIGN_OR_RAISE(bool same_location, + CatalogUtils::SameTableLocation(scope.table_path, context_->GetPath())); + const std::string read_branch = BranchManager::NormalizeBranch(context_->GetBranch()); + if (!same_location || scope.branch != read_branch) { + return Status::Invalid(fmt::format( + "The split was authorized for table {} (branch {}), but the read targets {} (branch " + "{}). Plan the scan for the table and branch being read.", + scope.table_path, scope.branch, context_->GetPath(), read_branch)); + } + const TableQueryAuthResult& rules = *split.GetAuthResult(); + const std::vector<DataField>& fields = table_schema_->Fields(); + // The field ids are checked first, so a renamed column is reported as such rather than as a + // missing one. + PAIMON_RETURN_NOT_OK(rules.ValidateFieldIds(scope.rule_field_ids, fields)); + PAIMON_RETURN_NOT_OK(rules.ValidateAgainstSchema(fields)); + // The columns the rules read themselves are the catalog's own and need no authorization. + for (const DataField& field : fields) { + const bool read = filter_fields.count(field.Name()) > 0 || !output_fields_ || + ObjectUtils::Contains(output_fields_.value(), field.Name()); + if (!read) { + continue; + } + auto authorized = scope.authorized_field_ids.find(field.Name()); + if (authorized == scope.authorized_field_ids.end()) { + return Status::Invalid( + fmt::format("Column '{}' is read, but the scan did not ask the catalog to " + "authorize it: name it in ScanContextBuilder::SetReadFieldNames().", + field.Name())); + } + if (authorized->second != field.Id()) { + return Status::Invalid(fmt::format( + "Column '{}' is read as field {}, but the scan authorized field {} under that " + "name: the schema changed since the plan. Plan the scan again.", + field.Name(), field.Id(), authorized->second)); + } + } + return Status::OK(); +} + +Result<std::unique_ptr<BatchReader>> QueryAuthTableRead::CreateReader( + const std::shared_ptr<Split>& split) { + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split); + if (auth_split == nullptr || auth_split->GetAuthResult() == nullptr) { + return Status::Invalid(fmt::format( + "A split of a table with '{}=true' must be planned by a scan of the table through its " + "catalog, which attaches the row filters and column masks it is read with.", + Options::QUERY_AUTH_ENABLED)); + } + const TableQueryAuthResult& rules = *auth_split->GetAuthResult(); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> filter, PostMaskFilter(rules)); + std::set<std::string> filter_fields; + if (filter != nullptr) { + PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(filter, &filter_fields)); + } + Status valid = ValidateSplit(*auth_split, filter_fields); + if (!valid.ok()) { + return valid.WithMessage(fmt::format("Query auth read of {} (schema {}): {}", + context_->GetPath(), table_schema_->Id(), + valid.message())); + } + + if (rules.RejectsAllRows()) { + return std::unique_ptr<BatchReader>(std::make_unique<ConcatBatchReader>( + std::vector<std::unique_ptr<BatchReader>>(), arrow_pool_)); + } + + // The columns read only for the rules or for the predicate, dropped from the output. + std::vector<std::string> extra_fields; + if (output_fields_) { + std::set<std::string> output(output_fields_->begin(), output_fields_->end()); + std::vector<std::string> visible = output_fields_.value(); + for (const std::string& field : filter_fields) { + if (output.count(field) == 0) { + visible.push_back(field); + } + } + std::set<std::string> required = rules.RequiredFields(visible); + required.insert(filter_fields.begin(), filter_fields.end()); + for (const DataField& field : table_schema_->Fields()) { + if (required.count(field.Name()) > 0 && output.count(field.Name()) == 0) { + extra_fields.push_back(field.Name()); + } + } + } + PAIMON_ASSIGN_OR_RAISE(TableRead * inner, GetOrCreateInnerRead(extra_fields)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, + inner->CreateReader(auth_split->GetSplit())); + // Wrapping it can fail, and `PredicateBatchReader::Create()` keeps the reader it refuses. + ScopeGuard close_guard([&reader]() { + if (reader) { + reader->Close(); + } + }); + + // The row filter sees the raw values, the predicate of the read the masked ones. Each reads + // its columns decoded, which `PredicateBatchReader` requires. Behind the row filter no column + // stays dictionary-encoded, including nested children: a dictionary would still hold the + // values of the rows it drops. + const std::map<std::string, std::shared_ptr<MaskTransform>> no_masks; + if (rules.RowFilter() != nullptr) { + reader = std::make_unique<QueryAuthBatchReader>( + std::move(reader), no_masks, + /*dropped_fields=*/std::set<std::string>(), rules.RowFilterFields(), arrow_pool_, + DataField::ConvertDataFieldsToArrowStructType(table_schema_->Fields())); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<PredicateBatchReader> filtered, + PredicateBatchReader::Create(std::move(reader), rules.RowFilter(), arrow_pool_)); + reader = std::make_unique<RowFilterBoundaryReader>(std::move(filtered)); + } + std::set<std::string> dropped_fields(extra_fields.begin(), extra_fields.end()); + if (!rules.Masks().empty() || !filter_fields.empty() || + (filter == nullptr && !dropped_fields.empty())) { + reader = std::make_unique<QueryAuthBatchReader>( + std::move(reader), rules.Masks(), + filter == nullptr ? dropped_fields : std::set<std::string>(), filter_fields, + arrow_pool_); + } + if (filter != nullptr) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<PredicateBatchReader> filtered, + PredicateBatchReader::Create(std::move(reader), filter, arrow_pool_)); + reader = std::move(filtered); + if (!dropped_fields.empty()) { + reader = std::make_unique<QueryAuthBatchReader>( + std::move(reader), no_masks, dropped_fields, + /*decoded_fields=*/std::set<std::string>(), arrow_pool_); + } + } + close_guard.Release(); + return reader; +} + +Result<std::unique_ptr<BatchReader>> QueryAuthTableRead::CreateReader( + const std::vector<std::shared_ptr<Split>>& splits) { + std::vector<std::unique_ptr<BatchReader>> readers; + readers.reserve(splits.size()); + ScopeGuard cleanup_guard([&readers]() { + for (const std::unique_ptr<BatchReader>& reader : readers) { + if (reader) { + reader->Close(); + } + } + }); + for (const std::shared_ptr<Split>& split : splits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, CreateReader(split)); + readers.push_back(std::move(reader)); + } + cleanup_guard.Release(); + return std::unique_ptr<BatchReader>( + std::make_unique<ConcatBatchReader>(std::move(readers), arrow_pool_)); +} + +} // namespace paimon
diff --git a/src/paimon/core/table/source/query_auth_table_read.h b/src/paimon/core/table/source/query_auth_table_read.h new file mode 100644 index 0000000..47c3962 --- /dev/null +++ b/src/paimon/core/table/source/query_auth_table_read.h
@@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <functional> +#include <map> +#include <memory> +#include <optional> +#include <set> +#include <string> +#include <vector> + +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/result.h" +#include "paimon/table/source/table_read.h" + +namespace arrow { +class MemoryPool; +} // namespace arrow + +namespace paimon { +class Metrics; +class Predicate; +class QueryAuthSplit; + +/// Applies the column masks to each batch, each computed from the raw batch, and drops the columns +/// only the rules needed, mirroring Java `TableQueryAuthResult.doAuth`. The masked columns, and +/// the columns a predicate evaluated next reads, are output decoded. +class QueryAuthBatchReader : public BatchReader { + public: + /// @param masks The masks to apply, keyed by the column they replace; a mask whose column is + /// not in a batch is not applied. + /// @param dropped_fields The columns removed from each batch after masking. + /// @param decoded_fields The unmasked columns output no longer dictionary-encoded or large + /// strings; other unmasked columns are output as read unless `logical_type` decodes + /// them. + /// @param logical_type When set, decodes every dictionary, including nested ones, to this + /// logical schema so a row filter cannot leave unauthorized values in a dictionary. + QueryAuthBatchReader(std::unique_ptr<BatchReader> reader, + std::map<std::string, std::shared_ptr<MaskTransform>> masks, + std::set<std::string> dropped_fields, std::set<std::string> decoded_fields, + std::shared_ptr<arrow::MemoryPool> arrow_pool, + std::shared_ptr<arrow::DataType> logical_type = nullptr); + + Result<BatchReader::ReadBatch> NextBatch() override; + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return reader_->GetReaderMetrics(); + } + + void Close() override { + reader_->Close(); + } + + private: + /// Returns `batch` with `decode_inputs_` no longer dictionary-encoded or large strings; the + /// other columns are left as they are. + Result<std::shared_ptr<arrow::StructArray>> DecodeInputs( + const std::shared_ptr<arrow::StructArray>& batch) const; + + std::unique_ptr<BatchReader> reader_; + std::map<std::string, std::shared_ptr<MaskTransform>> masks_; + std::set<std::string> dropped_fields_; + std::set<std::string> decoded_fields_; + /// `decoded_fields_` and the columns `masks_` read. + std::set<std::string> decode_inputs_; + std::shared_ptr<arrow::MemoryPool> arrow_pool_; + std::shared_ptr<arrow::DataType> logical_type_; +}; + +/// Reads a table with `query-auth.enabled=true`, mirroring Java `AbstractDataTableRead` and +/// `ReadTransform`: the rows of each `QueryAuthSplit` go through the row filter, the column masks, +/// then the read's own predicate, which is not pushed into the files. +class QueryAuthTableRead : public TableRead { + public: + /// Builds the read the rules are applied around from the context it is given. + using InnerReadFactory = + std::function<Result<std::unique_ptr<TableRead>>(std::unique_ptr<ReadContext>)>; + + /// @param context The read as the caller asked for it. + /// @param table_schema The schema the read reads with. + /// @param create_inner Builds the read the rules are applied around. + static Result<std::unique_ptr<TableRead>> Create( + const std::shared_ptr<ReadContext>& context, + const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, InnerReadFactory create_inner); + + Result<std::unique_ptr<BatchReader>> CreateReader(const std::shared_ptr<Split>& split) override; + Result<std::unique_ptr<BatchReader>> CreateReader( + const std::vector<std::shared_ptr<Split>>& splits) override; + + private: + QueryAuthTableRead(const std::shared_ptr<ReadContext>& context, + const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, + InnerReadFactory create_inner, + std::optional<std::vector<std::string>> output_fields); + + /// The predicate evaluated on the masked rows of a split read under `rules`: the whole + /// predicate with `EnablePredicateFilter(true)`, else its conjuncts on masked columns. + Result<std::shared_ptr<Predicate>> PostMaskFilter(const TableQueryAuthResult& rules) const; + + /// Checks that `split` was authorized for this table and branch, against the columns this + /// read sees and for the columns it outputs or filters on. `filter_fields` are the operands + /// of `PostMaskFilter()`. + Status ValidateSplit(const QueryAuthSplit& split, + const std::set<std::string>& filter_fields) const; + + /// The read of the requested columns plus `extra_fields`, built on first use. + Result<TableRead*> GetOrCreateInnerRead(const std::vector<std::string>& extra_fields); + + std::shared_ptr<ReadContext> context_; + std::shared_ptr<TableSchema> table_schema_; + std::shared_ptr<arrow::MemoryPool> arrow_pool_; + InnerReadFactory create_inner_; + /// The table columns the caller reads, or null when it reads them all. + std::optional<std::vector<std::string>> output_fields_; + std::map<std::vector<std::string>, std::unique_ptr<TableRead>> inner_reads_; +}; + +} // namespace paimon
diff --git a/src/paimon/core/table/source/query_auth_table_read_test.cpp b/src/paimon/core/table/source/query_auth_table_read_test.cpp new file mode 100644 index 0000000..32e39d7 --- /dev/null +++ b/src/paimon/core/table/source/query_auth_table_read_test.cpp
@@ -0,0 +1,399 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/table/source/query_auth_table_read.h" + +#include <cstdint> +#include <functional> +#include <map> +#include <memory> +#include <set> +#include <string> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" +#include "gtest/gtest.h" +#include "paimon/common/reader/predicate_batch_reader.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/core/table/source/query_auth_split.h" +#include "paimon/defs.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/result.h" +#include "paimon/status.h" +#include "paimon/table/source/split.h" +#include "paimon/table/source/table_read.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +namespace { + +class OneBatchReader : public BatchReader { + public: + explicit OneBatchReader(std::shared_ptr<arrow::Array> batch) : batch_(std::move(batch)) {} + + Result<ReadBatch> NextBatch() override { + if (batch_ == nullptr) { + return BatchReader::MakeEofBatch(); + } + auto c_array = std::make_unique<::ArrowArray>(); + auto c_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch_, c_array.get(), c_schema.get())); + batch_.reset(); + return std::make_pair(std::move(c_array), std::move(c_schema)); + } + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return nullptr; + } + + void Close() override {} + + private: + std::shared_ptr<arrow::Array> batch_; +}; + +class CloseCountingReader : public BatchReader { + public: + explicit CloseCountingReader(int32_t* closes) : closes_(closes) {} + + Result<ReadBatch> NextBatch() override { + return BatchReader::MakeEofBatch(); + } + + std::shared_ptr<Metrics> GetReaderMetrics() const override { + return nullptr; + } + + void Close() override { + ++*closes_; + } + + private: + int32_t* closes_; +}; + +class DummySplit final : public Split {}; + +class SplitTableRead : public TableRead { + public: + using ReaderFactory = + std::function<Result<std::unique_ptr<BatchReader>>(const std::shared_ptr<Split>&)>; + + explicit SplitTableRead(ReaderFactory create_reader) + : create_reader_(std::move(create_reader)) {} + + Result<std::unique_ptr<BatchReader>> CreateReader( + const std::shared_ptr<Split>& split) override { + return create_reader_(split); + } + + Result<std::unique_ptr<BatchReader>> CreateReader( + const std::vector<std::shared_ptr<Split>>& /*splits*/) override { + return Status::NotImplemented("the query auth read opens one split at a time"); + } + + private: + ReaderFactory create_reader_; +}; + +QueryAuthTableRead::InnerReadFactory CreateSplitTableRead( + SplitTableRead::ReaderFactory create_reader) { + return [create_reader = std::move(create_reader)]( + std::unique_ptr<ReadContext>) -> Result<std::unique_ptr<TableRead>> { + return std::unique_ptr<TableRead>(std::make_unique<SplitTableRead>(create_reader)); + }; +} + +} // namespace + +TEST(QueryAuthTableReadTest, TestBatchReaderDecodesDictionaryAndLargeStringColumns) { + auto json = [](const std::shared_ptr<arrow::DataType>& type, const std::string& values) { + return arrow::ipc::internal::json::ArrayFromJSON(type, values).ValueOrDie(); + }; + const auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + std::shared_ptr<arrow::Array> name = + arrow::DictionaryArray::FromArrays(dictionary_type, json(arrow::int32(), "[0, 1, null]"), + json(arrow::utf8(), R"(["alice", "bob"])")) + .ValueOrDie(); + std::shared_ptr<arrow::Array> city = + arrow::DictionaryArray::FromArrays(dictionary_type, json(arrow::int32(), "[0, 0, 0]"), + json(arrow::utf8(), R"(["x"])")) + .ValueOrDie(); + std::shared_ptr<arrow::Array> nick = + arrow::DictionaryArray::FromArrays(dictionary_type, json(arrow::int32(), "[1, null, 0]"), + json(arrow::utf8(), R"(["al", "bo"])")) + .ValueOrDie(); + std::shared_ptr<arrow::Array> phone = + json(arrow::large_utf8(), R"(["13800001111", "13900002222", null])"); + std::shared_ptr<arrow::Array> batch = + arrow::StructArray::Make( + {name, phone, city, nick, city}, + {arrow::field("name", dictionary_type), arrow::field("phone", arrow::large_utf8()), + arrow::field("city", dictionary_type), arrow::field("nick", dictionary_type), + arrow::field("alias", dictionary_type)}) + .ValueOrDie(); + const std::string name_ref = R"({"index": 0, "name": "name", "type": "STRING"})"; + const std::string phone_ref = R"({"index": 1, "name": "phone", "type": "STRING"})"; + const std::string nick_ref = R"({"index": 3, "name": "nick", "type": "STRING"})"; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> rules, + TableQueryAuthResult::Create( + {}, + {{"name", fmt::format(R"({{"name": "SUBSTRING", "inputs": [{}, 1, 1]}})", name_ref)}, + {"phone", fmt::format(R"({{"name": "CONCAT", "inputs": [{}, "-"]}})", phone_ref)}, + {"city", R"({"name": "CONCAT", "inputs": ["***"]})"}, + {"alias", fmt::format(R"({{"name": "FIELD_REF", "fieldRef": {}}})", nick_ref)}})); + + QueryAuthBatchReader reader(std::make_unique<OneBatchReader>(batch), rules->Masks(), + /*dropped_fields=*/{}, /*decoded_fields=*/{}, + GetArrowPool(GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch masked_batch, reader.NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(masked_batch)); + std::shared_ptr<arrow::Array> masked = + arrow::ImportArray(masked_batch.first.get(), masked_batch.second.get()).ValueOrDie(); + ASSERT_TRUE(masked->ValidateFull().ok()) << masked->ValidateFull().ToString(); + auto masked_struct = checked_pointer_cast<arrow::StructArray>(masked); + for (const auto& [column, values] : std::vector<std::pair<std::string, std::string>>{ + {"name", R"(["a", "b", null])"}, + {"phone", R"(["13800001111-", "13900002222-", null])"}, + {"city", R"(["***", "***", "***"])"}, + {"alias", R"(["bo", null, "al"])"}}) { + std::shared_ptr<arrow::Array> field = masked_struct->GetFieldByName(column); + ASSERT_NE(nullptr, field) << column; + ASSERT_TRUE(field->Equals(*json(arrow::utf8(), values))) + << column << ": " << field->ToString(); + } + ASSERT_TRUE(masked_struct->GetFieldByName("nick")->type()->Equals(*dictionary_type)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch eof, reader.NextBatch()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<Predicate> name_or_phone, + PredicateBuilder::Or( + {PredicateBuilder::Equal(/*field_index=*/0, "name", FieldType::STRING, + Literal(FieldType::STRING, "bob", 3)), + PredicateBuilder::Equal(/*field_index=*/1, "phone", FieldType::STRING, + Literal(FieldType::STRING, "13800001111", 11))})); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<PredicateBatchReader> filtered, + PredicateBatchReader::Create(std::make_unique<QueryAuthBatchReader>( + std::make_unique<OneBatchReader>(batch), + std::map<std::string, std::shared_ptr<MaskTransform>>(), + /*dropped_fields=*/std::set<std::string>(), + /*decoded_fields=*/std::set<std::string>{"name", "phone"}, + GetArrowPool(GetDefaultPool())), + name_or_phone, GetArrowPool(GetDefaultPool()))); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch filtered_batch, filtered->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(filtered_batch)); + auto filtered_struct = checked_pointer_cast<arrow::StructArray>( + arrow::ImportArray(filtered_batch.first.get(), filtered_batch.second.get()).ValueOrDie()); + ASSERT_TRUE(filtered_struct->GetFieldByName("name")->Equals( + *json(arrow::utf8(), R"(["alice", "bob"])"))); + ASSERT_TRUE(filtered_struct->GetFieldByName("phone")->Equals( + *json(arrow::utf8(), R"(["13800001111", "13900002222"])"))); + ASSERT_TRUE(filtered_struct->GetFieldByName("city")->type()->Equals(*dictionary_type)); + + ASSERT_OK_AND_ASSIGN( + filtered, + PredicateBatchReader::Create( + std::make_unique<QueryAuthBatchReader>( + std::make_unique<OneBatchReader>(batch), + std::map<std::string, std::shared_ptr<MaskTransform>>(), + /*dropped_fields=*/std::set<std::string>(), + /*decoded_fields=*/std::set<std::string>{"name"}, GetArrowPool(GetDefaultPool()), + /*logical_type=*/ + arrow::struct_( + {arrow::field("name", arrow::utf8()), arrow::field("phone", arrow::utf8()), + arrow::field("city", arrow::utf8()), arrow::field("nick", arrow::utf8()), + arrow::field("alias", arrow::utf8())})), + PredicateBuilder::Equal(/*field_index=*/0, "name", FieldType::STRING, + Literal(FieldType::STRING, "bob", 3)), + GetArrowPool(GetDefaultPool()))); + ASSERT_OK_AND_ASSIGN(filtered_batch, filtered->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(filtered_batch)); + filtered_struct = checked_pointer_cast<arrow::StructArray>( + arrow::ImportArray(filtered_batch.first.get(), filtered_batch.second.get()).ValueOrDie()); + ASSERT_TRUE(filtered_struct->GetFieldByName("nick")->Equals(*json(arrow::utf8(), "[null]"))); + ASSERT_TRUE(filtered_struct->GetFieldByName("city")->Equals(*json(arrow::utf8(), R"(["x"])"))); + ASSERT_TRUE(filtered_struct->GetFieldByName("phone")->type()->Equals(*arrow::large_utf8())); +} + +TEST(QueryAuthTableReadTest, TestClosesTheReadersOfSplitsBeforeOneThatFails) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TableSchema> created, + TableSchema::Create(/*schema_id=*/0, arrow::schema({arrow::field("id", arrow::int32())}), + /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::QUERY_AUTH_ENABLED, "true"}})); + std::shared_ptr<TableSchema> table_schema = std::move(created); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> context, + ReadContextBuilder("table_root_path").Finish()); + int32_t closes = 0; + auto failing_split = std::make_shared<DummySplit>(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TableRead> read, + QueryAuthTableRead::Create( + std::shared_ptr<ReadContext>(std::move(context)), table_schema, + GetArrowPool(GetDefaultPool()), + CreateSplitTableRead([&closes, failing_split](const std::shared_ptr<Split>& split) + -> Result<std::unique_ptr<BatchReader>> { + if (split == failing_split) { + return Status::IOError("cannot open the split"); + } + return std::unique_ptr<BatchReader>(std::make_unique<CloseCountingReader>(&closes)); + }))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> no_rules, + TableQueryAuthResult::Create({}, {})); + QueryAuthScope scope; + scope.table_path = "table_root_path"; + scope.branch = "main"; + for (const DataField& field : table_schema->Fields()) { + scope.authorized_field_ids.emplace(field.Name(), field.Id()); + } + auto authorized = [&](const std::shared_ptr<Split>& split) -> std::shared_ptr<Split> { + return std::make_shared<QueryAuthSplit>(split, no_rules, scope); + }; + + ASSERT_NOK_WITH_MSG(read->CreateReader(std::vector<std::shared_ptr<Split>>{ + authorized(std::make_shared<DummySplit>()), authorized(failing_split)}), + "cannot open the split"); + ASSERT_EQ(1, closes); + ASSERT_NOK_WITH_MSG( + read->CreateReader(std::vector<std::shared_ptr<Split>>{ + authorized(std::make_shared<DummySplit>()), authorized(std::make_shared<DummySplit>()), + std::make_shared<DummySplit>()}), + "must be planned by a scan of the table through its catalog"); + ASSERT_EQ(3, closes); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader, + read->CreateReader(std::vector<std::shared_ptr<Split>>{ + authorized(std::make_shared<DummySplit>()), + authorized(std::make_shared<DummySplit>())})); + ASSERT_EQ(3, closes); + reader->Close(); + ASSERT_EQ(5, closes); +} + +TEST(QueryAuthTableReadTest, TestRowFilterRemovesRowsFromTheBatchWithBitmap) { + auto json = [](const std::shared_ptr<arrow::DataType>& type, const std::string& values) { + return arrow::ipc::internal::json::ArrayFromJSON(type, values).ValueOrDie(); + }; + const arrow::FieldVector fields = { + arrow::field("tenant", arrow::utf8()), + arrow::field("value", arrow::int32()), + arrow::field("email", arrow::utf8()), + arrow::field("profile", arrow::struct_({arrow::field("email", arrow::utf8())})), + arrow::field("emails", arrow::list(arrow::utf8())), + arrow::field("contacts", arrow::map(arrow::utf8(), arrow::utf8()))}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> created, + TableSchema::Create(/*schema_id=*/0, arrow::schema(fields), + /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::QUERY_AUTH_ENABLED, "true"}})); + std::shared_ptr<TableSchema> table_schema = std::move(created); + // The dictionary contains B's email even after every remaining index belongs to A. + std::shared_ptr<arrow::Array> email = + arrow::DictionaryArray::FromArrays( + arrow::dictionary(arrow::int32(), arrow::utf8()), json(arrow::int32(), "[0, 1, 2]"), + json(arrow::utf8(), R"(["a@example", "b@example", "a2@example"])")) + .ValueOrDie(); + std::shared_ptr<arrow::Array> offsets = json(arrow::int32(), "[0, 1, 2, 3]"); + std::shared_ptr<arrow::Array> profile = + arrow::StructArray::Make({email}, std::vector<std::string>{"email"}).ValueOrDie(); + std::shared_ptr<arrow::Array> emails = + arrow::ListArray::FromArrays(*offsets, *email).ValueOrDie(); + std::shared_ptr<arrow::Array> contacts = + arrow::MapArray::FromArrays(offsets, json(arrow::utf8(), R"(["home", "home", "home"])"), + email) + .ValueOrDie(); + std::shared_ptr<arrow::Array> batch = + arrow::StructArray::Make( + {json(arrow::utf8(), R"(["A", "B", "A"])"), json(arrow::int32(), "[1, 2, 3]"), email, + profile, emails, contacts}, + std::vector<std::string>{"tenant", "value", "email", "profile", "emails", "contacts"}) + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> rules, + TableQueryAuthResult::Create( + {R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0,)" + R"( "name": "tenant", "type": "STRING"}}, "function": "EQUAL", "literals": ["A"]})"}, + {})); + QueryAuthScope scope; + scope.table_path = "table_root_path"; + scope.branch = "main"; + for (const DataField& field : table_schema->Fields()) { + scope.authorized_field_ids.emplace(field.Name(), field.Id()); + } + scope.rule_field_ids = rules->BindFieldIds(table_schema->Fields()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> context, + ReadContextBuilder("table_root_path").Finish()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TableRead> read, + QueryAuthTableRead::Create( + std::shared_ptr<ReadContext>(std::move(context)), table_schema, + GetArrowPool(GetDefaultPool()), + CreateSplitTableRead( + [batch](const std::shared_ptr<Split>&) -> Result<std::unique_ptr<BatchReader>> { + return std::unique_ptr<BatchReader>(std::make_unique<OneBatchReader>(batch)); + }))); + auto split = std::make_shared<QueryAuthSplit>(std::make_shared<DummySplit>(), rules, scope); + + auto assert_only_tenant_a = [&](std::unique_ptr<BatchReader> reader) { + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch_with_bitmap)); + auto& [read_batch, bitmap] = batch_with_bitmap; + std::function<void(const ArrowArray&)> assert_no_dictionary = [&](const ArrowArray& array) { + ASSERT_EQ(nullptr, array.dictionary); + for (int64_t i = 0; i < array.n_children; ++i) { + assert_no_dictionary(*array.children[i]); + } + }; + assert_no_dictionary(*read_batch.first); + std::shared_ptr<arrow::Array> array = + arrow::ImportArray(read_batch.first.get(), read_batch.second.get()).ValueOrDie(); + auto rows = checked_pointer_cast<arrow::StructArray>(array); + ASSERT_TRUE(rows->ValidateFull().ok()); + ASSERT_TRUE(rows->Equals( + *json(arrow::struct_(fields), + R"([["A", 1, "a@example", ["a@example"], ["a@example"], [["home", "a@example"]]], + ["A", 3, "a2@example", ["a2@example"], ["a2@example"], [["home", "a2@example"]]]])"))) + << rows->ToString(); + ASSERT_EQ(2, bitmap.Cardinality()); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap eof, reader->NextBatchWithBitmap()); + ASSERT_TRUE(BatchReader::IsEofBatch(eof)); + reader->Close(); + }; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader, read->CreateReader(split)); + assert_only_tenant_a(std::move(reader)); + ASSERT_OK_AND_ASSIGN(reader, read->CreateReader(std::vector<std::shared_ptr<Split>>{split})); + assert_only_tenant_a(std::move(reader)); +} + +} // namespace paimon::test
diff --git a/src/paimon/core/table/source/query_auth_table_scan.cpp b/src/paimon/core/table/source/query_auth_table_scan.cpp new file mode 100644 index 0000000..e1c72b2 --- /dev/null +++ b/src/paimon/core/table/source/query_auth_table_scan.cpp
@@ -0,0 +1,383 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/core/table/source/query_auth_table_scan.h" + +#include <algorithm> +#include <chrono> +#include <map> +#include <optional> +#include <utility> +#include <vector> + +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/object_utils.h" +#include "paimon/core/catalog/catalog_utils.h" +#include "paimon/core/catalog/query_auth_catalog.h" +#include "paimon/core/schema/schema_manager.h" +#include "paimon/core/table/source/plan_impl.h" +#include "paimon/core/table/source/query_auth_split.h" +#include "paimon/core/utils/branch_manager.h" +#include "paimon/defs.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/schema/schema.h" +#include "paimon/status.h" + +namespace paimon { +namespace { + +std::set<std::string> Intersect(const std::set<std::string>& left, + const std::set<std::string>& right) { + std::set<std::string> result; + for (const std::string& value : left) { + if (right.count(value) > 0) { + result.insert(value); + } + } + return result; +} + +Status CheckNoMaskedPartitionFilter(const std::set<std::string>& partition_filter_fields, + const std::set<std::string>& masked_fields) { + const std::set<std::string> masked = Intersect(partition_filter_fields, masked_fields); + if (!masked.empty()) { + return Status::NotImplemented(fmt::format( + "A partition filter cannot be enforced on masked partition key '{}': it is matched " + "against the raw partition value. Filter on the column through the predicate instead.", + *masked.begin())); + } + return Status::OK(); +} + +const DataField* FindField(const std::vector<DataField>& fields, const std::string& name) { + auto iter = std::find_if(fields.begin(), fields.end(), + [&](const DataField& field) { return field.Name() == name; }); + return iter != fields.end() ? &*iter : nullptr; +} + +std::string FieldIdOrAbsent(const DataField* field) { + return field != nullptr ? std::to_string(field->Id()) : std::string("absent"); +} + +} // namespace + +Result<std::unique_ptr<TableScan>> QueryAuthTableScan::Create( + const std::shared_ptr<ScanContext>& context, const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<FileSystem>& file_system, const std::string& branch, + InnerScanFactory create_inner) { + if (context->GetCatalog() == nullptr || !context->GetIdentifier()) { + return Status::NotImplemented(fmt::format( + "Cannot read a table with '{}=true' without the catalog that serves its row filters " + "and column masks: build the scan with ScanContextBuilder::WithCatalog().", + Options::QUERY_AUTH_ENABLED)); + } + const QueryAuthCatalog* auth_catalog = AsQueryAuthCatalog(context->GetCatalog()); + if (auth_catalog == nullptr) { + return Status::NotImplemented(fmt::format( + "Cannot read a table with '{}=true': its catalog does not support the query auth " + "interface.", + Options::QUERY_AUTH_ENABLED)); + } + if (context->GetGlobalIndexResult() != nullptr) { + return Status::NotImplemented( + "Search is not supported on a query-auth table: the index ranks raw values, which a " + "column mask invalidates."); + } + if (context->GetRealtimeContext() != nullptr) { + return Status::NotImplemented("Real-time reads are not supported on a query-auth table."); + } + // The rules are asked for the branch planned, even when only the options name it. + PAIMON_ASSIGN_OR_RAISE(std::string resolved_branch, + BranchManager::ResolveBranch(context->GetIdentifier(), + context->GetOptions(), branch, "scan")); + PAIMON_ASSIGN_OR_RAISE( + Identifier auth_identifier, + CatalogUtils::ResolvedBranchIdentifier(context->GetIdentifier().value(), resolved_branch)); + PAIMON_RETURN_NOT_OK(CatalogUtils::CheckPathIsTableLocation( + auth_identifier, context->GetCatalog()->GetTableLocation(auth_identifier), + context->GetPath(), "scan")); + std::set<std::string> user_filter_fields; + std::set<std::string> partition_filter_fields; + const std::shared_ptr<ScanFilter>& scan_filter = context->GetScanFilters(); + if (scan_filter != nullptr) { + if (scan_filter->GetPredicate() != nullptr) { + PAIMON_RETURN_NOT_OK( + PredicateUtils::GetAllNames(scan_filter->GetPredicate(), &user_filter_fields)); + } + for (const auto& partition : scan_filter->GetPartitionFilters()) { + for (const auto& [key, value] : partition) { + partition_filter_fields.insert(key); + } + } + } + // As in Java, the operands of the filters are authorized along with the projected columns. + const std::optional<std::vector<std::string>>& read_field_names = context->GetReadFieldNames(); + std::vector<std::string> requested = read_field_names.value_or(std::vector<std::string>()); + requested.insert(requested.end(), user_filter_fields.begin(), user_filter_fields.end()); + requested.insert(requested.end(), partition_filter_fields.begin(), + partition_filter_fields.end()); + std::vector<std::string> columns; + std::set<std::string> seen; + for (const std::string& name : requested) { + if (!table_schema->GetField(name).ok()) { + return Status::NotImplemented(fmt::format( + "Column '{}' is not a column of the table schema, which is all a query-auth scan " + "can authorize.", + name)); + } + if (seen.insert(name).second) { + columns.push_back(name); + } + } + std::optional<std::vector<std::string>> authorization_columns; + if (read_field_names) { + authorization_columns = std::move(columns); + } + return std::unique_ptr<TableScan>(new QueryAuthTableScan( + context, auth_catalog, std::move(auth_identifier), table_schema, file_system, + resolved_branch, std::move(create_inner), std::move(authorization_columns), + std::move(user_filter_fields), std::move(partition_filter_fields))); +} + +QueryAuthTableScan::QueryAuthTableScan( + const std::shared_ptr<ScanContext>& context, const QueryAuthCatalog* auth_catalog, + Identifier auth_identifier, const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<FileSystem>& file_system, const std::string& branch, + InnerScanFactory create_inner, std::optional<std::vector<std::string>> authorization_columns, + std::set<std::string> user_filter_fields, std::set<std::string> partition_filter_fields) + : context_(context), + auth_catalog_(auth_catalog), + auth_identifier_(std::move(auth_identifier)), + table_schema_(table_schema), + file_system_(file_system), + branch_(branch), + create_inner_(std::move(create_inner)), + authorization_columns_(std::move(authorization_columns)), + user_filter_fields_(std::move(user_filter_fields)), + partition_filter_fields_(std::move(partition_filter_fields)) {} + +std::shared_ptr<Metrics> QueryAuthTableScan::GetMetrics() const { + return inner_ != nullptr ? inner_->GetMetrics() : TableScan::GetMetrics(); +} + +Status QueryAuthTableScan::AddErrorContext(const Status& status) const { + return status.WithMessage(fmt::format("Query auth of {} (branch {}, schema {}): {}", + auth_identifier_.ToString(), branch_, table_schema_->Id(), + status.message())); +} + +Status QueryAuthTableScan::ValidateAgainstCurrentSchemas(const TableQueryAuthResult& rules) const { + SchemaManager schema_manager(file_system_, context_->GetPath(), branch_); + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest, + schema_manager.Latest()); + const std::shared_ptr<TableSchema>& latest_schema = latest ? latest.value() : table_schema_; + const std::vector<DataField>& planned_fields = table_schema_->Fields(); + const std::vector<DataField>& latest_fields = latest_schema->Fields(); + // The catalog authorizes the columns by their names in the latest schema, so each must still + // be the field the scan plans under that name, whether or not the rules mention it. + std::vector<std::string> columns; + if (authorization_columns_) { + columns = authorization_columns_.value(); + } else { + for (const DataField& field : planned_fields) { + columns.push_back(field.Name()); + } + } + auto validate_columns = [&](const std::vector<DataField>& fields, int64_t schema_id, + const char* source) -> Status { + for (const std::string& column : columns) { + const DataField* planned = FindField(planned_fields, column); + const DataField* current = FindField(fields, column); + if (planned == nullptr || current == nullptr || planned->Id() != current->Id()) { + return Status::Invalid(fmt::format( + "Column '{}' is field {} in schema {} the scan plans with, but field {} in {} " + "{}, under which the catalog authorized it. Create the scan again.", + column, FieldIdOrAbsent(planned), table_schema_->Id(), FieldIdOrAbsent(current), + source, schema_id)); + } + } + return Status::OK(); + }; + PAIMON_RETURN_NOT_OK(validate_columns(latest_fields, latest_schema->Id(), "the latest schema")); + PAIMON_RETURN_NOT_OK(rules.ValidateAgainstSchema(latest_fields)); + PAIMON_RETURN_NOT_OK(rules.ValidateFieldIds(rules.BindFieldIds(latest_fields), planned_fields)); + // A rule column the planned schema does not hold would leave every split unreadable. + PAIMON_RETURN_NOT_OK( + rules.ValidateFieldIds(rules.BindFieldIds(planned_fields), planned_fields)); + + // The catalog can advance its schema before publishing it under the table path. It + // authorizes current names, so binding those names only to local field ids could expose an + // old column that the catalog has dropped and replaced with another of the same name. + std::string location; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Schema> catalog_schema, + auth_catalog_->LoadTableSchemaAndLocation(auth_identifier_, &location)); + // The identifier may have been bound to another table since the scan was created. + PAIMON_RETURN_NOT_OK(CatalogUtils::CheckPathIsTableLocation(auth_identifier_, location, + context_->GetPath(), "scan")); + auto catalog_data_schema = std::dynamic_pointer_cast<DataSchema>(catalog_schema); + if (catalog_data_schema == nullptr) { + return Status::Invalid("Query auth requires a data table schema from the catalog."); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> c_schema, + catalog_data_schema->GetArrowSchema()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> arrow_schema, + arrow::ImportSchema(c_schema.get())); + PAIMON_ASSIGN_OR_RAISE(std::vector<DataField> catalog_fields, + DataField::ConvertArrowSchemaToDataFields(arrow_schema)); + PAIMON_RETURN_NOT_OK( + validate_columns(catalog_fields, catalog_data_schema->Id(), "catalog schema")); + PAIMON_RETURN_NOT_OK(rules.ValidateAgainstSchema(catalog_fields)); + return rules.ValidateFieldIds(rules.BindFieldIds(catalog_fields), planned_fields); +} + +Result<std::shared_ptr<TableQueryAuthResult>> QueryAuthTableScan::LoadAndValidateRules() const { + const auto start = std::chrono::steady_clock::now(); + Result<std::shared_ptr<TableQueryAuthResult>> rules = + auth_catalog_->AuthTableQuery(auth_identifier_, authorization_columns_); + Status status = rules.status(); + const char* failed_step = "catalog request"; + if (status.ok()) { + status = ValidateAgainstCurrentSchemas(*rules.value()); + failed_step = "validation"; + } + if (logger_->IsLevelEnabled(PAIMON_LOG_LEVEL_DEBUG)) { + const std::string outcome = + status.ok() ? fmt::format("{} filters, {} masks", rules.value()->Filter().size(), + rules.value()->ColumnMasking().size()) + : fmt::format("{} failed with {}", failed_step, status.CodeAsString()); + const int64_t elapsed_us = std::chrono::duration_cast<std::chrono::microseconds>( + std::chrono::steady_clock::now() - start) + .count(); + const std::string message = fmt::format( + "query auth of {} (branch {}, schema {}) took {} us: {}", auth_identifier_.ToString(), + branch_, table_schema_->Id(), elapsed_us, outcome); + PAIMON_LOG_DEBUG(logger_, "%s", message.c_str()); + } + if (!status.ok()) { + return AddErrorContext(status); + } + return rules; +} + +Result<std::unique_ptr<TableScan>> QueryAuthTableScan::CreateInnerScan( + const TableQueryAuthResult& rules) const { + const std::set<std::string> masked_fields = rules.MaskedFields(); + PAIMON_RETURN_NOT_OK(CheckNoMaskedPartitionFilter(partition_filter_fields_, masked_fields)); + const std::shared_ptr<ScanFilter>& scan_filter = context_->GetScanFilters(); + std::vector<std::shared_ptr<Predicate>> pushed; + // Raw statistics of a masked column must not prune: the read evaluates those conjuncts on the + // masked values, so they drop rows only then and a pushed limit could stop short of them. + const bool masked_conjuncts_excluded = !Intersect(user_filter_fields_, masked_fields).empty(); + if (scan_filter != nullptr && scan_filter->GetPredicate() != nullptr) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<Predicate> unmasked, + PredicateUtils::ExcludePredicateWithFields(scan_filter->GetPredicate(), masked_fields)); + if (unmasked != nullptr) { + pushed.push_back(unmasked); + } + } + // The row filter only removes what the read would drop anyway, so it can prune. + if (rules.RowFilter() != nullptr) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Predicate> row_filter, + PredicateUtils::CreatePickedFieldFilter( + rules.RowFilter(), ObjectUtils::CreateIdentifierToIndexMap( + table_schema_->FieldNames()))); + if (row_filter != nullptr) { + pushed.push_back(row_filter); + } + } + std::shared_ptr<Predicate> pushed_predicate; + if (!pushed.empty()) { + PAIMON_ASSIGN_OR_RAISE(pushed_predicate, PredicateBuilder::And(pushed)); + } + auto pushed_filter = std::make_shared<ScanFilter>( + pushed_predicate, + scan_filter != nullptr ? scan_filter->GetPartitionFilters() + : std::vector<std::map<std::string, std::string>>(), + scan_filter != nullptr ? scan_filter->GetBucketFilter() : std::nullopt); + // As in Java, a row filter on data columns turns the limit off too: it is pushed down as a + // predicate, which the inner scan does not push a limit past. + auto inner_context = std::make_shared<ScanContext>( + context_->GetPath(), context_->IsStreamingMode(), + masked_conjuncts_excluded ? std::optional<int32_t>() : context_->GetLimit(), pushed_filter, + /*global_index_result=*/nullptr, /*realtime_context=*/nullptr, context_->GetMemoryPool(), + context_->GetExecutor(), context_->GetSpecificFileSystem(), + context_->GetSpecificTableSchema(), context_->GetOptions(), context_->GetCache(), + context_->GetFormatTable(), context_->GetCatalog(), context_->GetIdentifier(), + context_->GetReadFieldNames()); + return create_inner_(inner_context); +} + +Status QueryAuthTableScan::ValidatePushdownCompatibility(const TableQueryAuthResult& rules) const { + const std::set<std::string> masked_fields = rules.MaskedFields(); + PAIMON_RETURN_NOT_OK(CheckNoMaskedPartitionFilter(partition_filter_fields_, masked_fields)); + // A pushed row filter that changed would prune the rows the new one lets through, and a + // column masked since the predicate on it was pushed would have pruned by its raw values. + if (!rules.SameFilter(*pushed_rules_) || + !ObjectUtils::ContainsAll(pushed_rules_->MaskedFields(), + Intersect(user_filter_fields_, masked_fields))) { + return Status::Invalid( + "Query auth rules changed since the scan pushed down its filters. Recreate the scan " + "to apply the new rules."); + } + return Status::OK(); +} + +Result<std::shared_ptr<Plan>> QueryAuthTableScan::CreatePlan() { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableQueryAuthResult> rules, LoadAndValidateRules()); + if (inner_ == nullptr) { + Result<std::unique_ptr<TableScan>> inner = CreateInnerScan(*rules); + if (!inner.ok()) { + return AddErrorContext(inner.status()); + } + inner_ = std::move(inner).value(); + pushed_rules_ = rules; + } else { + Status compatible = ValidatePushdownCompatibility(*rules); + if (!compatible.ok()) { + return AddErrorContext(compatible); + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, inner_->CreatePlan()); + // Bound to the schema planned, so the read can tell a column since renamed or replaced, and + // one added since, from the columns the rules were authorized for. + QueryAuthScope scope; + scope.table_path = context_->GetPath(); + scope.branch = branch_; + for (const DataField& field : table_schema_->Fields()) { + if (!authorization_columns_ || + ObjectUtils::Contains(authorization_columns_.value(), field.Name())) { + scope.authorized_field_ids.emplace(field.Name(), field.Id()); + } + } + scope.rule_field_ids = rules->BindFieldIds(table_schema_->Fields()); + // Every split carries the rules, none or not: the read refuses a split planned without them. + std::vector<std::shared_ptr<Split>> splits; + splits.reserve(plan->Splits().size()); + for (const std::shared_ptr<Split>& split : plan->Splits()) { + splits.push_back(std::make_shared<QueryAuthSplit>(split, rules, scope)); + } + return std::shared_ptr<Plan>(std::make_shared<PlanImpl>(plan->SnapshotId(), splits)); +} + +} // namespace paimon
diff --git a/src/paimon/core/table/source/query_auth_table_scan.h b/src/paimon/core/table/source/query_auth_table_scan.h new file mode 100644 index 0000000..0f7d9a1 --- /dev/null +++ b/src/paimon/core/table/source/query_auth_table_scan.h
@@ -0,0 +1,113 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include <functional> +#include <memory> +#include <optional> +#include <set> +#include <string> +#include <vector> + +#include "paimon/catalog/catalog.h" +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/core/schema/table_schema.h" +#include "paimon/fs/file_system.h" +#include "paimon/logging.h" +#include "paimon/result.h" +#include "paimon/scan_context.h" +#include "paimon/table/source/plan.h" +#include "paimon/table/source/table_scan.h" + +namespace paimon { +class Metrics; +class QueryAuthCatalog; + +/// Plans a table with `query-auth.enabled=true`, mirroring Java `AbstractDataTableScan.authQuery`: +/// every plan loads and checks the rules and wraps each split in a `QueryAuthSplit`. Only the row +/// filter and the predicate conjuncts on unmasked columns are pushed down, and the limit only when +/// the predicate has no conjunct on a masked column. +class QueryAuthTableScan : public TableScan { + public: + /// Builds the scan the rules are applied around from the context it is given. + using InnerScanFactory = + std::function<Result<std::unique_ptr<TableScan>>(const std::shared_ptr<ScanContext>&)>; + + /// @param context The scan as the caller asked for it; it must name the catalog and the table. + static Result<std::unique_ptr<TableScan>> Create( + const std::shared_ptr<ScanContext>& context, + const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<FileSystem>& file_system, const std::string& branch, + InnerScanFactory create_inner); + + Result<std::shared_ptr<Plan>> CreatePlan() override; + + std::shared_ptr<Metrics> GetMetrics() const override; + + private: + QueryAuthTableScan(const std::shared_ptr<ScanContext>& context, + const QueryAuthCatalog* auth_catalog, Identifier auth_identifier, + const std::shared_ptr<TableSchema>& table_schema, + const std::shared_ptr<FileSystem>& file_system, const std::string& branch, + InnerScanFactory create_inner, + std::optional<std::vector<std::string>> authorization_columns, + std::set<std::string> user_filter_fields, + std::set<std::string> partition_filter_fields); + + /// Adds the table, branch and schema of this scan to the message of `status`. + Status AddErrorContext(const Status& status) const; + + /// Checks both the latest published schema and the catalog's current schema. Fails when a + /// column the catalog is asked to authorize is another field than the scan plans with, when + /// `rules` do not fit either schema, when they reference a column the planned schema does + /// not hold or holds as another field, or when the catalog no longer keeps the table at the + /// scan path. + Status ValidateAgainstCurrentSchemas(const TableQueryAuthResult& rules) const; + + /// Asks the catalog for the rules of this query and checks them with + /// `ValidateAgainstCurrentSchemas()`. + Result<std::shared_ptr<TableQueryAuthResult>> LoadAndValidateRules() const; + + /// Builds the scan the rules are applied around, with what `rules` let it push down. + Result<std::unique_ptr<TableScan>> CreateInnerScan(const TableQueryAuthResult& rules) const; + + /// Fails when `rules` would need pruning the inner scan already pushed down to be undone. + Status ValidatePushdownCompatibility(const TableQueryAuthResult& rules) const; + + /// Holds the catalog `auth_catalog_` points to. + std::shared_ptr<ScanContext> context_; + const QueryAuthCatalog* auth_catalog_; + /// The table the rules are asked for, naming the branch planned. + Identifier auth_identifier_; + std::shared_ptr<TableSchema> table_schema_; + std::shared_ptr<FileSystem> file_system_; + std::string branch_; + InnerScanFactory create_inner_; + /// The columns the catalog is asked to authorize: the projection plus the filter operands, or + /// null for every column. + std::optional<std::vector<std::string>> authorization_columns_; + std::set<std::string> user_filter_fields_; + std::set<std::string> partition_filter_fields_; + /// Built by the first plan, from the rules it loaded. + std::unique_ptr<TableScan> inner_; + std::shared_ptr<TableQueryAuthResult> pushed_rules_; + std::shared_ptr<Logger> logger_ = Logger::GetLogger("QueryAuthTableScan"); +}; + +} // namespace paimon
diff --git a/src/paimon/core/table/source/split.cpp b/src/paimon/core/table/source/split.cpp index e1ff8a1..52cc07d 100644 --- a/src/paimon/core/table/source/split.cpp +++ b/src/paimon/core/table/source/split.cpp
@@ -17,6 +17,9 @@ * under the License. */ +#include <map> +#include <optional> +#include <string> #include <utility> #include "fmt/format.h" @@ -31,6 +34,7 @@ #include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/deletion_file.h" #include "paimon/core/table/source/fallback_data_split.h" +#include "paimon/core/table/source/query_auth_split.h" #include "paimon/core/utils/object_serializer.h" #include "paimon/global_index/indexed_split.h" #include "paimon/io/byte_array_input_stream.h" @@ -42,6 +46,27 @@ namespace paimon { struct DataFileMeta; namespace { + +/// Checks a serialized count before allocating: each entry needs at least `min_entry_bytes`. +Result<int32_t> ReadSplitCount(DataInputStream* in, int64_t min_entry_bytes, + const char* split_type) { + PAIMON_ASSIGN_OR_RAISE(int32_t count, in->ReadValue<int32_t>()); + PAIMON_ASSIGN_OR_RAISE(int64_t pos, in->GetPos()); + PAIMON_ASSIGN_OR_RAISE(int64_t stream_length, in->Length()); + if (count < 0 || count * min_entry_bytes > stream_length - pos) { + return Status::Invalid( + fmt::format("invalid {}, {} entries do not fit the split", split_type, count)); + } + return count; +} + +/// Java DataInput.readBoolean() treats any nonzero byte as true. Reading into a C++ bool would +/// create an invalid representation for a byte other than 0 or 1. +Result<bool> ReadSplitFlag(DataInputStream* in) { + PAIMON_ASSIGN_OR_RAISE(int8_t value, in->ReadValue<int8_t>()); + return value != 0; +} + Status WriteDataSplit(const std::shared_ptr<DataSplitImpl>& data_split_impl, MemorySegmentOutputStream* out, const std::shared_ptr<MemoryPool>& pool) { out->WriteValue<int64_t>(DataSplitImpl::MAGIC); @@ -93,7 +118,7 @@ std::optional<int32_t> total_buckets; if (version >= 6) { - PAIMON_ASSIGN_OR_RAISE(bool total_buckets_exist, in->ReadValue<bool>()); + PAIMON_ASSIGN_OR_RAISE(bool total_buckets_exist, ReadSplitFlag(in)); if (total_buckets_exist) { PAIMON_ASSIGN_OR_RAISE(total_buckets, in->ReadValue<int32_t>()); } @@ -115,9 +140,9 @@ PAIMON_ASSIGN_OR_RAISE(data_deletion_files, DeletionFile::DeserializeList(in, version)); bool is_streaming = false; - PAIMON_ASSIGN_OR_RAISE(is_streaming, in->ReadValue<bool>()); + PAIMON_ASSIGN_OR_RAISE(is_streaming, ReadSplitFlag(in)); bool raw_convertible = false; - PAIMON_ASSIGN_OR_RAISE(raw_convertible, in->ReadValue<bool>()); + PAIMON_ASSIGN_OR_RAISE(raw_convertible, ReadSplitFlag(in)); DataSplitImpl::Builder builder(partition, bucket, bucket_path, std::move(data_files)); builder.WithTotalBuckets(total_buckets) @@ -134,39 +159,284 @@ return builder.Build(); } +Status WriteIndexedSplit(const std::shared_ptr<IndexedSplitImpl>& indexed_split_impl, + MemorySegmentOutputStream* out, const std::shared_ptr<MemoryPool>& pool) { + out->WriteValue<int64_t>(IndexedSplitImpl::MAGIC); + out->WriteValue<int32_t>(IndexedSplitImpl::VERSION); + auto inner_split_impl = + std::dynamic_pointer_cast<DataSplitImpl>(indexed_split_impl->GetDataSplit()); + if (!inner_split_impl) { + return Status::Invalid("inner split in IndexedSplit is supposed to be DataSplit"); + } + PAIMON_RETURN_NOT_OK(WriteDataSplit(inner_split_impl, out, pool)); + auto row_ranges = indexed_split_impl->RowRanges(); + out->WriteValue<int32_t>(row_ranges.size()); + for (const auto& range : row_ranges) { + out->WriteValue<int64_t>(range.from); + out->WriteValue<int64_t>(range.to); + } + + auto scores = indexed_split_impl->Scores(); + if (!scores.empty()) { + out->WriteValue<bool>(true); + out->WriteValue<int32_t>(scores.size()); + for (float score : scores) { + out->WriteValue<float>(CanonicalizeFloatingPoint(score)); + } + } else { + out->WriteValue<bool>(false); + } + return Status::OK(); +} + +/// Reads an IndexedSplit after its magic number. +Result<std::shared_ptr<IndexedSplitImpl>> ReadIndexedSplitWithoutMagicNumber( + DataInputStream* in, const std::shared_ptr<MemoryPool>& pool) { + PAIMON_ASSIGN_OR_RAISE(int32_t version, in->ReadValue<int32_t>()); + if (version != IndexedSplitImpl::VERSION) { + return Status::Invalid(fmt::format("Unsupported IndexedSplit version: {}", version)); + } + PAIMON_ASSIGN_OR_RAISE(int64_t data_split_magic, in->ReadValue<int64_t>()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataSplitImpl> data_split, + ReadDataSplitWithoutMagicNumber(data_split_magic, in, pool)); + PAIMON_ASSIGN_OR_RAISE(int32_t range_size, + ReadSplitCount(in, 2 * sizeof(int64_t), "IndexedSplit")); + std::vector<Range> row_ranges; + row_ranges.reserve(range_size); + for (int32_t i = 0; i < range_size; ++i) { + PAIMON_ASSIGN_OR_RAISE(int64_t range_from, in->ReadValue<int64_t>()); + PAIMON_ASSIGN_OR_RAISE(int64_t range_to, in->ReadValue<int64_t>()); + row_ranges.emplace_back(range_from, range_to); + } + std::vector<float> scores; + PAIMON_ASSIGN_OR_RAISE(bool has_scores, ReadSplitFlag(in)); + if (has_scores) { + PAIMON_ASSIGN_OR_RAISE(int32_t scores_length, + ReadSplitCount(in, sizeof(float), "IndexedSplit")); + scores.resize(scores_length); + for (int32_t i = 0; i < scores_length; ++i) { + PAIMON_ASSIGN_OR_RAISE(float score, in->ReadValue<float>()); + scores[i] = score; + } + } + return std::make_shared<IndexedSplitImpl>(data_split, row_ranges, scores); +} + +/// Java `QueryAuthSplit.writeString` of a non-null string: its length, then the UTF-8 bytes. +void WriteAuthString(const std::string& value, MemorySegmentOutputStream* out) { + out->WriteValue<int32_t>(static_cast<int32_t>(value.size())); + out->Write(value.data(), static_cast<uint32_t>(value.size())); +} + +/// A null string reads as an empty one, which the rules refuse. +Result<std::string> ReadAuthString(DataInputStream* in) { + PAIMON_ASSIGN_OR_RAISE(int32_t length, in->ReadValue<int32_t>()); + if (length < 0) { + return std::string(); + } + PAIMON_ASSIGN_OR_RAISE(int64_t pos, in->GetPos()); + PAIMON_ASSIGN_OR_RAISE(int64_t stream_length, in->Length()); + if (length > stream_length - pos) { + return Status::Invalid( + fmt::format("invalid QueryAuthSplit, a string of {} bytes overruns the split", length)); + } + std::string value(length, '\0'); + PAIMON_RETURN_NOT_OK(in->Read(value.data(), length)); + return value; +} + +Status WriteWithSplitSerializer(const std::shared_ptr<Split>& split, MemorySegmentOutputStream* out, + const std::shared_ptr<MemoryPool>& pool) { + out->WriteValue<int64_t>(QueryAuthSplit::kSplitSerializerMagic); + out->WriteValue<int32_t>(QueryAuthSplit::kSplitSerializerVersion); + if (auto data_split_impl = std::dynamic_pointer_cast<DataSplitImpl>(split)) { + out->WriteValue<int32_t>(QueryAuthSplit::kSplitSerializerDataSplit); + return WriteDataSplit(data_split_impl, out, pool); + } + if (auto indexed_split_impl = std::dynamic_pointer_cast<IndexedSplitImpl>(split)) { + out->WriteValue<int32_t>(QueryAuthSplit::kSplitSerializerIndexedSplit); + return WriteIndexedSplit(indexed_split_impl, out, pool); + } + return Status::Invalid("a QueryAuthSplit can only wrap a DataSplit or an IndexedSplit"); +} + +/// Java `QueryAuthSplit.serialize`, framed by `SplitSerializer`. +Status WriteQueryAuthSplitWithSplitSerializer(const QueryAuthSplit& split, + MemorySegmentOutputStream* out, + const std::shared_ptr<MemoryPool>& pool) { + out->WriteValue<int64_t>(QueryAuthSplit::kSplitSerializerMagic); + out->WriteValue<int32_t>(QueryAuthSplit::kSplitSerializerVersion); + out->WriteValue<int32_t>(QueryAuthSplit::kSplitSerializerQueryAuthSplit); + PAIMON_RETURN_NOT_OK(WriteWithSplitSerializer(split.GetSplit(), out, pool)); + const std::shared_ptr<TableQueryAuthResult>& auth_result = split.GetAuthResult(); + if (auth_result == nullptr) { + return Status::Invalid("a QueryAuthSplit must carry the rules it was planned with"); + } + out->WriteValue<bool>(true); + out->WriteValue<bool>(true); + out->WriteValue<int32_t>(static_cast<int32_t>(auth_result->Filter().size())); + for (const std::string& filter : auth_result->Filter()) { + WriteAuthString(filter, out); + } + out->WriteValue<bool>(true); + out->WriteValue<int32_t>(static_cast<int32_t>(auth_result->ColumnMasking().size())); + for (const auto& [column, transform] : auth_result->ColumnMasking()) { + WriteAuthString(column, out); + WriteAuthString(transform, out); + } + return Status::OK(); +} + +void WriteFieldIds(const std::map<std::string, int32_t>& field_ids, + MemorySegmentOutputStream* out) { + out->WriteValue<int32_t>(static_cast<int32_t>(field_ids.size())); + for (const auto& [field, id] : field_ids) { + WriteAuthString(field, out); + out->WriteValue<int32_t>(id); + } +} + +/// A split without a scope, as Java wrote it, is written in the Java format alone. +Status WriteQueryAuthSplit(const QueryAuthSplit& split, MemorySegmentOutputStream* out, + const std::shared_ptr<MemoryPool>& pool) { + const std::optional<QueryAuthScope>& scope = split.GetScope(); + if (!scope) { + return WriteQueryAuthSplitWithSplitSerializer(split, out, pool); + } + out->WriteValue<int64_t>(QueryAuthSplit::kScopedMagic); + out->WriteValue<int32_t>(QueryAuthSplit::kScopedVersion); + PAIMON_RETURN_NOT_OK(WriteQueryAuthSplitWithSplitSerializer(split, out, pool)); + WriteFieldIds(scope->authorized_field_ids, out); + WriteFieldIds(scope->rule_field_ids, out); + WriteAuthString(scope->table_path, out); + WriteAuthString(scope->branch, out); + return Status::OK(); +} + +/// Reads a split Java `SplitSerializer` wrote, after its magic number; `in_query_auth_split` when +/// it is the split a `QueryAuthSplit` wraps. +Result<std::shared_ptr<Split>> ReadWithSplitSerializerWithoutMagicNumber( + DataInputStream* in, const std::shared_ptr<MemoryPool>& pool, bool in_query_auth_split) { + PAIMON_ASSIGN_OR_RAISE(int32_t version, in->ReadValue<int32_t>()); + if (version != QueryAuthSplit::kSplitSerializerVersion) { + return Status::Invalid(fmt::format("Unsupported split serializer version: {}", version)); + } + PAIMON_ASSIGN_OR_RAISE(int32_t type, in->ReadValue<int32_t>()); + if (type == QueryAuthSplit::kSplitSerializerDataSplit) { + PAIMON_ASSIGN_OR_RAISE(int64_t magic, in->ReadValue<int64_t>()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataSplitImpl> data_split, + ReadDataSplitWithoutMagicNumber(magic, in, pool)); + return std::shared_ptr<Split>(data_split); + } + if (type == QueryAuthSplit::kSplitSerializerIndexedSplit) { + PAIMON_ASSIGN_OR_RAISE(int64_t magic, in->ReadValue<int64_t>()); + if (magic != IndexedSplitImpl::MAGIC) { + return Status::Invalid("invalid IndexedSplit, wrong magic number"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<IndexedSplitImpl> indexed_split, + ReadIndexedSplitWithoutMagicNumber(in, pool)); + return std::shared_ptr<Split>(indexed_split); + } + if (type != QueryAuthSplit::kSplitSerializerQueryAuthSplit) { + return Status::Invalid(fmt::format("Unsupported split type: {}", type)); + } + // Refused before reading it, so nested wrappers cannot recurse without bound. + if (in_query_auth_split) { + return Status::Invalid("invalid QueryAuthSplit, it wraps another QueryAuthSplit"); + } + PAIMON_ASSIGN_OR_RAISE(int64_t inner_magic, in->ReadValue<int64_t>()); + if (inner_magic != QueryAuthSplit::kSplitSerializerMagic) { + return Status::Invalid( + "invalid QueryAuthSplit, its split is not framed by SplitSerializer"); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<Split> inner, + ReadWithSplitSerializerWithoutMagicNumber(in, pool, /*in_query_auth_split=*/true)); + PAIMON_ASSIGN_OR_RAISE(bool has_auth_result, ReadSplitFlag(in)); + if (!has_auth_result) { + // Rules that cannot be read must not read as no rules. + return Status::Invalid("invalid QueryAuthSplit, it carries no rules"); + } + std::vector<std::string> filter; + PAIMON_ASSIGN_OR_RAISE(bool has_filter, ReadSplitFlag(in)); + if (has_filter) { + PAIMON_ASSIGN_OR_RAISE(int32_t size, ReadSplitCount(in, sizeof(int32_t), "QueryAuthSplit")); + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::string value, ReadAuthString(in)); + filter.push_back(std::move(value)); + } + } + std::map<std::string, std::string> column_masking; + PAIMON_ASSIGN_OR_RAISE(bool has_column_masking, ReadSplitFlag(in)); + if (has_column_masking) { + PAIMON_ASSIGN_OR_RAISE(int32_t size, + ReadSplitCount(in, 2 * sizeof(int32_t), "QueryAuthSplit")); + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::string column, ReadAuthString(in)); + PAIMON_ASSIGN_OR_RAISE(std::string transform, ReadAuthString(in)); + if (!column_masking.emplace(std::move(column), std::move(transform)).second) { + return Status::Invalid("invalid QueryAuthSplit, its rules mask a column twice"); + } + } + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableQueryAuthResult> auth_result, + TableQueryAuthResult::Create(filter, column_masking)); + return std::shared_ptr<Split>( + std::make_shared<QueryAuthSplit>(inner, auth_result, /*scope=*/std::nullopt)); +} + +Result<std::map<std::string, int32_t>> ReadFieldIds(DataInputStream* in) { + PAIMON_ASSIGN_OR_RAISE(int32_t size, ReadSplitCount(in, 2 * sizeof(int32_t), "QueryAuthSplit")); + std::map<std::string, int32_t> field_ids; + for (int32_t i = 0; i < size; ++i) { + PAIMON_ASSIGN_OR_RAISE(std::string field, ReadAuthString(in)); + PAIMON_ASSIGN_OR_RAISE(int32_t id, in->ReadValue<int32_t>()); + if (!field_ids.emplace(std::move(field), id).second) { + return Status::Invalid("invalid QueryAuthSplit, its scope names a column twice"); + } + } + return field_ids; +} + +/// Reads a split `WriteQueryAuthSplit` wrote with its scope, after the magic number. +Result<std::shared_ptr<Split>> ReadScopedQueryAuthSplitWithoutMagicNumber( + DataInputStream* in, const std::shared_ptr<MemoryPool>& pool) { + PAIMON_ASSIGN_OR_RAISE(int32_t version, in->ReadValue<int32_t>()); + if (version != QueryAuthSplit::kScopedVersion) { + return Status::Invalid(fmt::format("Unsupported QueryAuthSplit version: {}", version)); + } + PAIMON_ASSIGN_OR_RAISE(int64_t magic, in->ReadValue<int64_t>()); + if (magic != QueryAuthSplit::kSplitSerializerMagic) { + return Status::Invalid( + "invalid QueryAuthSplit, its split is not framed by SplitSerializer"); + } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr<Split> split, + ReadWithSplitSerializerWithoutMagicNumber(in, pool, /*in_query_auth_split=*/false)); + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split); + if (auth_split == nullptr) { + return Status::Invalid("invalid QueryAuthSplit, its envelope holds another split"); + } + QueryAuthScope scope; + PAIMON_ASSIGN_OR_RAISE(scope.authorized_field_ids, ReadFieldIds(in)); + PAIMON_ASSIGN_OR_RAISE(scope.rule_field_ids, ReadFieldIds(in)); + PAIMON_ASSIGN_OR_RAISE(scope.table_path, ReadAuthString(in)); + PAIMON_ASSIGN_OR_RAISE(scope.branch, ReadAuthString(in)); + return std::shared_ptr<Split>(std::make_shared<QueryAuthSplit>( + auth_split->GetSplit(), auth_split->GetAuthResult(), std::move(scope))); +} + } // namespace Result<std::string> Split::Serialize(const std::shared_ptr<Split>& split, const std::shared_ptr<MemoryPool>& pool) { MemorySegmentOutputStream out(MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool); - if (auto data_split_impl = std::dynamic_pointer_cast<DataSplitImpl>(split)) { + if (auto query_auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(split)) { + PAIMON_RETURN_NOT_OK(WriteQueryAuthSplit(*query_auth_split, &out, pool)); + } else if (auto data_split_impl = std::dynamic_pointer_cast<DataSplitImpl>(split)) { PAIMON_RETURN_NOT_OK(WriteDataSplit(data_split_impl, &out, pool)); } else if (auto indexed_split_impl = std::dynamic_pointer_cast<IndexedSplitImpl>(split)) { - out.WriteValue<int64_t>(IndexedSplitImpl::MAGIC); - out.WriteValue<int32_t>(IndexedSplitImpl::VERSION); - auto inner_split_impl = - std::dynamic_pointer_cast<DataSplitImpl>(indexed_split_impl->GetDataSplit()); - if (!inner_split_impl) { - return Status::Invalid("inner split in IndexedSplit is supposed to be DataSplit"); - } - PAIMON_RETURN_NOT_OK(WriteDataSplit(inner_split_impl, &out, pool)); - auto row_ranges = indexed_split_impl->RowRanges(); - out.WriteValue<int32_t>(row_ranges.size()); - for (const auto& range : row_ranges) { - out.WriteValue<int64_t>(range.from); - out.WriteValue<int64_t>(range.to); - } - - auto scores = indexed_split_impl->Scores(); - if (!scores.empty()) { - out.WriteValue<bool>(true); - out.WriteValue<int32_t>(scores.size()); - for (float score : scores) { - out.WriteValue<float>(CanonicalizeFloatingPoint(score)); - } - } else { - out.WriteValue<bool>(false); - } + PAIMON_RETURN_NOT_OK(WriteIndexedSplit(indexed_split_impl, &out, pool)); } else if (std::dynamic_pointer_cast<FormatDataSplit>(split) != nullptr) { // A format table's plan has no cross-runtime encoding, and one of paimon-cpp's own would // let a plan made here be handed to a runtime that cannot read it back. @@ -189,37 +459,30 @@ int64_t magic = -1; PAIMON_ASSIGN_OR_RAISE(magic, in.ReadValue<int64_t>()); - if (magic == IndexedSplitImpl::MAGIC) { - PAIMON_ASSIGN_OR_RAISE(int32_t version, in.ReadValue<int32_t>()); - if (version != IndexedSplitImpl::VERSION) { - return Status::Invalid(fmt::format("Unsupported IndexedSplit version: {}", version)); + if (magic == QueryAuthSplit::kSplitSerializerMagic || magic == QueryAuthSplit::kScopedMagic) { + std::shared_ptr<Split> split; + if (magic == QueryAuthSplit::kScopedMagic) { + PAIMON_ASSIGN_OR_RAISE(split, ReadScopedQueryAuthSplitWithoutMagicNumber(&in, pool)); + } else { + PAIMON_ASSIGN_OR_RAISE(split, ReadWithSplitSerializerWithoutMagicNumber( + &in, pool, /*in_query_auth_split=*/false)); } - PAIMON_ASSIGN_OR_RAISE(int64_t data_split_magic, in.ReadValue<int64_t>()); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataSplitImpl> data_split, - ReadDataSplitWithoutMagicNumber(data_split_magic, &in, pool)); - PAIMON_ASSIGN_OR_RAISE(int32_t range_size, in.ReadValue<int32_t>()); - std::vector<Range> row_ranges; - row_ranges.reserve(range_size); - for (int32_t i = 0; i < range_size; ++i) { - PAIMON_ASSIGN_OR_RAISE(int64_t range_from, in.ReadValue<int64_t>()); - PAIMON_ASSIGN_OR_RAISE(int64_t range_to, in.ReadValue<int64_t>()); - row_ranges.emplace_back(range_from, range_to); + PAIMON_ASSIGN_OR_RAISE(int64_t pos, in.GetPos()); + PAIMON_ASSIGN_OR_RAISE(int64_t stream_length, in.Length()); + if (pos != stream_length) { + return Status::Invalid( + fmt::format("invalid split byte stream, remaining {} bytes after deserializing", + stream_length - pos)); } - std::vector<float> scores; - PAIMON_ASSIGN_OR_RAISE(bool has_scores, in.ReadValue<bool>()); - if (has_scores) { - PAIMON_ASSIGN_OR_RAISE(int32_t scores_length, in.ReadValue<int32_t>()); - scores.resize(scores_length); - for (int32_t i = 0; i < scores_length; ++i) { - PAIMON_ASSIGN_OR_RAISE(float score, in.ReadValue<float>()); - scores[i] = score; - } - } + return split; + } else if (magic == IndexedSplitImpl::MAGIC) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<IndexedSplitImpl> indexed_split, + ReadIndexedSplitWithoutMagicNumber(&in, pool)); // TODO(lisizhuo.lsz): support fallback split in IndexedSplit PAIMON_ASSIGN_OR_RAISE(int64_t pos, in.GetPos()); PAIMON_ASSIGN_OR_RAISE(int64_t stream_length, in.Length()); if (pos == stream_length) { - return std::make_shared<IndexedSplitImpl>(data_split, row_ranges, scores); + return indexed_split; } else if (pos == stream_length - 1) { return Status::Invalid( "invalid IndexedSplit, do not support FallbackSplit in IndexedSplit"); @@ -236,7 +499,7 @@ if (pos == stream_length) { return data_split; } else if (pos == stream_length - 1) { - PAIMON_ASSIGN_OR_RAISE(bool is_fallback, in.ReadValue<bool>()); + PAIMON_ASSIGN_OR_RAISE(bool is_fallback, ReadSplitFlag(&in)); return std::make_shared<FallbackDataSplit>(data_split, is_fallback); } else { return Status::Invalid(fmt::format(
diff --git a/src/paimon/core/table/system/global_system_tables.cpp b/src/paimon/core/table/system/global_system_tables.cpp index 647da76..809c53a 100644 --- a/src/paimon/core/table/system/global_system_tables.cpp +++ b/src/paimon/core/table/system/global_system_tables.cpp
@@ -37,6 +37,7 @@ #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/sensitive_config_utils.h" #include "paimon/common/utils/string_utils.h" +#include "paimon/core/catalog/query_auth_catalog.h" #include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_entry.h" @@ -45,6 +46,7 @@ #include "paimon/core/manifest/manifest_file.h" #include "paimon/core/manifest/manifest_file_meta.h" #include "paimon/core/manifest/manifest_list.h" +#include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/snapshot.h" #include "paimon/core/utils/branch_manager.h" @@ -130,12 +132,14 @@ std::string database_name; std::string table_name; std::shared_ptr<DataSchema> schema; + std::optional<std::string> location; }; // Match Java CatalogUtils::listAllTables: tolerate databases or tables removed concurrently, but // propagate all other catalog and schema errors instead of returning incomplete system-table rows. Result<std::vector<CatalogTableInfo>> LoadAllDataTables(const Catalog& catalog) { std::vector<CatalogTableInfo> result; + const auto* auth_catalog = dynamic_cast<const QueryAuthCatalog*>(&catalog); PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> databases, catalog.ListDatabases()); for (const std::string& database : databases) { Result<std::vector<std::string>> tables_result = catalog.ListTables(database); @@ -147,7 +151,13 @@ } for (const std::string& table : tables_result.value()) { Identifier identifier(database, table); - Result<std::shared_ptr<Schema>> schema_result = catalog.LoadTableSchema(identifier); + // Keep the query auth setting and location from one response, so sys.partitions + // cannot apply an old table's disabled setting to a recreated table's files. + std::optional<std::string> location; + Result<std::shared_ptr<Schema>> schema_result = + auth_catalog != nullptr + ? auth_catalog->LoadTableSchemaAndLocation(identifier, &location.emplace()) + : catalog.LoadTableSchema(identifier); if (!schema_result.ok()) { if (schema_result.status().IsNotExist()) { continue; @@ -160,7 +170,7 @@ return Status::Invalid("catalog returned a non-data schema for ", identifier.ToString()); } - result.push_back({database, table, std::move(data_schema)}); + result.push_back({database, table, std::move(data_schema), std::move(location)}); } } return result; @@ -179,6 +189,20 @@ std::map<std::string, FileStats> by_partition; }; +// As for a data read, the latest published schema can enable query auth before the catalog's +// schema does. +Result<bool> LatestSchemaEnablesQueryAuth(const std::shared_ptr<FileSystem>& fs, + const std::string& table_path) { + SchemaManager schema_manager(fs, table_path); + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest_schema, + schema_manager.Latest()); + if (!latest_schema) { + return false; + } + return OptionsUtils::GetValueFromMap<bool>(latest_schema.value()->Options(), + Options::QUERY_AUTH_ENABLED, false); +} + // Read the latest snapshot's data files and aggregate statistics. Result<AggregatedFileStats> AggregateFileStats(const std::shared_ptr<FileSystem>& fs, const std::string& table_path, @@ -504,10 +528,19 @@ if (data_schema->PartitionKeys().empty()) { continue; } + // The partition values of a query-auth table are what its masks and row filter may hide. + PAIMON_ASSIGN_OR_RAISE(bool query_auth, + OptionsUtils::GetValueFromMap<bool>( + data_schema->Options(), Options::QUERY_AUTH_ENABLED, false)); + if (query_auth) { + continue; + } // Match Java's toAllPartitions by ignoring only concurrent table deletion. All other // metadata and I/O errors must fail the query. - Result<std::string> table_path_result = context_.catalog->GetTableLocation(id); + Result<std::string> table_path_result = table.location + ? Result<std::string>(table.location.value()) + : context_.catalog->GetTableLocation(id); if (!table_path_result.ok()) { if (table_path_result.status().IsNotExist()) { continue; @@ -526,6 +559,18 @@ return table_fs_result.status(); } + Result<bool> latest_query_auth_result = + LatestSchemaEnablesQueryAuth(table_fs_result.value(), table_path_result.value()); + if (!latest_query_auth_result.ok()) { + if (latest_query_auth_result.status().IsNotExist()) { + continue; + } + return latest_query_auth_result.status(); + } + if (latest_query_auth_result.value()) { + continue; + } + Result<AggregatedFileStats> file_stats_result = AggregateFileStats(table_fs_result.value(), table_path_result.value(), *data_schema); if (!file_stats_result.ok()) {
diff --git a/src/paimon/core/table/system/system_table.cpp b/src/paimon/core/table/system/system_table.cpp index 5a1a1de..6b8d04a 100644 --- a/src/paimon/core/table/system/system_table.cpp +++ b/src/paimon/core/table/system/system_table.cpp
@@ -26,11 +26,13 @@ #include <utility> #include <vector> +#include "fmt/format.h" #include "paimon/catalog/catalog.h" #include "paimon/catalog/identifier.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/schema/schema_manager.h" +#include "paimon/core/schema/schema_validation.h" #include "paimon/core/schema/table_schema.h" #include "paimon/core/table/system/audit_log_system_table.h" #include "paimon/core/table/system/binlog_system_table.h" @@ -174,6 +176,26 @@ return std::nullopt; } +// The system tables that read the rows or the file statistics of a query-auth table are refused +// as soon as they are loaded: `$files` and `$binlog` as in Java, `$audit_log` and `$ro`, whose +// reads would not apply the rules, and `$manifests`, whose partition statistics hold raw +// partition values that a mask or a row filter hides. The other system tables report only +// metadata. +Status ValidateQueryAuthSystemTable(const std::string& system_table_name, + const std::shared_ptr<TableSchema>& table_schema, + const std::map<std::string, std::string>& dynamic_options) { + std::string normalized_name = StringUtils::ToLowerCase(system_table_name); + if (table_schema == nullptr || (normalized_name != FilesSystemTable::kName && + normalized_name != BinlogSystemTable::kName && + normalized_name != AuditLogSystemTable::kName && + normalized_name != ReadOptimizedSystemTable::kName && + normalized_name != ManifestsSystemTable::kName)) { + return Status::OK(); + } + return SchemaValidation::RejectQueryAuth(*table_schema, dynamic_options, + fmt::format("System table '{}'", normalized_name)); +} + } // namespace bool SystemTableLoader::IsSupported(const std::string& system_table_name) { @@ -186,6 +208,8 @@ const std::map<std::string, std::string>& dynamic_options) { std::optional<SystemTableFactory> factory = FindSystemTableFactory(system_table_name); if (factory) { + PAIMON_RETURN_NOT_OK( + ValidateQueryAuthSystemTable(system_table_name, table_schema, dynamic_options)); return factory.value()(fs, table_path, table_schema, dynamic_options); } return Status::NotImplemented("unsupported system table: ", system_table_name); @@ -244,16 +268,18 @@ return GlobalSystemTableLoader::Load(parsed.system_table_name, context); } - SchemaManager schema_manager(fs, parsed.table_path, - parsed.branch.value_or(BranchManager::DEFAULT_MAIN_BRANCH)); + PAIMON_ASSIGN_OR_RAISE(std::string branch, BranchManager::ResolveBranch( + /*identifier=*/std::nullopt, dynamic_options, + parsed.branch, "system table read")); + SchemaManager schema_manager(fs, parsed.table_path, branch); PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest_schema, schema_manager.Latest()); if (!latest_schema) { return Status::NotExist("base table schema not found for system table path: ", path); } auto options = dynamic_options; - if (parsed.branch) { - options[Options::BRANCH] = parsed.branch.value(); + if (parsed.branch || options.count(Options::BRANCH) > 0) { + options[Options::BRANCH] = branch; } return Load(parsed.system_table_name, fs, parsed.table_path, latest_schema.value(), options); }
diff --git a/src/paimon/core/table/table.cpp b/src/paimon/core/table/table.cpp index ec131ef..5d53636 100644 --- a/src/paimon/core/table/table.cpp +++ b/src/paimon/core/table/table.cpp
@@ -34,6 +34,7 @@ #include "paimon/common/global_index/btree/btree_defs.h" #include "paimon/common/predicate/predicate_validator.h" #include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/binary_row_partition_computer.h" #include "paimon/common/utils/fields_comparator.h" @@ -42,6 +43,7 @@ #include "paimon/core/catalog/catalog_snapshot_commit.h" #include "paimon/core/catalog/catalog_utils.h" #include "paimon/core/catalog/commit_lock.h" +#include "paimon/core/catalog/query_auth_catalog.h" #include "paimon/core/catalog/renaming_snapshot_commit.h" #include "paimon/core/catalog/snapshot_commit.h" #include "paimon/core/catalog/version_managed_catalog.h" @@ -91,6 +93,8 @@ #include "paimon/core/table/source/key_value_table_read.h" #include "paimon/core/table/source/merge_tree_split_generator.h" #include "paimon/core/table/source/primary_key_index_batch_scan.h" +#include "paimon/core/table/source/query_auth_table_read.h" +#include "paimon/core/table/source/query_auth_table_scan.h" #include "paimon/core/table/source/read_optimized_scan_options.h" #include "paimon/core/table/source/realtime_table_scan.h" #include "paimon/core/table/source/snapshot/snapshot_reader.h" @@ -640,10 +644,108 @@ } }; +/// Returns `options` with `query-auth.enabled` turned on when the current settings of `branch` +/// enable it: the options of the schema the catalog holds for that branch, as Java loads the table +/// of the branch it reads, or, when the caller hands in a cached schema, those of the latest schema +/// of the branch under `path`, so a stale cache cannot skip the rules. A system table path naming +/// a branch, as `tbl$branch_dev$files`, reads that branch, so its settings are the ones looked up. +/// @param operation The operation reading `path`, as "scan", named in the error messages. +Result<std::map<std::string, std::string>> ResolveQueryAuthReadOptions( + const std::shared_ptr<Catalog>& catalog, const std::optional<Identifier>& identifier, + const std::map<std::string, std::string>& options, const std::string& branch, + const std::shared_ptr<FileSystem>& file_system, const std::string& path, + bool has_specific_table_schema, const std::string& operation) { + PAIMON_ASSIGN_OR_RAISE(std::optional<SystemTablePath> system_table_path, + SystemTableLoader::TryParsePath(path)); + std::string read_branch = branch; + std::string table_path = path; + if (system_table_path && !system_table_path->is_global) { + table_path = system_table_path->table_path; + if (system_table_path->branch) { + // Resolved as `SystemTableLoader::LoadFromPath()` resolves the branch it loads. + std::map<std::string, std::string> named_branch; + if (!BranchManager::IsMainBranch(branch)) { + named_branch[Options::BRANCH] = branch; + } + PAIMON_ASSIGN_OR_RAISE(read_branch, + BranchManager::ResolveBranch( + /*identifier=*/std::nullopt, named_branch, + system_table_path->branch, "system table read")); + } + } + bool enabled = false; + if (catalog != nullptr) { + PAIMON_RETURN_NOT_OK(BranchManager::CheckCatalogAddressableBranch(read_branch)); + if (!identifier) { + return Status::Invalid("a catalog read requires a table identifier"); + } + if (CatalogUtils::IsSystemDatabase(identifier.value().GetDatabaseName())) { + // Only a global system table has no data table policy to load from the catalog. + if (!system_table_path || !system_table_path->is_global) { + return Status::Invalid(fmt::format( + "Catalog identifier {} names a global system table, but the {} targets {}. " + "Use the identifier of the data table being read.", + identifier.value().ToString(), operation, path)); + } + } else { + PAIMON_ASSIGN_OR_RAISE(Identifier table_identifier, + CatalogUtils::BranchIdentifier(identifier.value(), read_branch)); + // A catalog that serves query auth hands the location back with the schema, so the + // two cannot come from different generations of a recreated table. + const QueryAuthCatalog* auth_catalog = AsQueryAuthCatalog(catalog); + std::shared_ptr<Schema> schema; + std::optional<std::string> location; + if (auth_catalog != nullptr) { + PAIMON_ASSIGN_OR_RAISE(schema, auth_catalog->LoadTableSchemaAndLocation( + table_identifier, &location.emplace())); + } else { + PAIMON_ASSIGN_OR_RAISE(schema, catalog->LoadTableSchema(table_identifier)); + } + std::shared_ptr<DataSchema> data_schema = std::dynamic_pointer_cast<DataSchema>(schema); + if (data_schema == nullptr) { + return Status::Invalid( + fmt::format("the catalog did not hand back a data table schema for {}", + table_identifier.ToString())); + } + PAIMON_ASSIGN_OR_RAISE( + enabled, SchemaValidation::QueryAuthReadEnabled(data_schema->Options(), options)); + // An enabled setting needs no check here: the query auth scan checks the location and + // the read takes only its splits. A disabled one would leave another table unguarded. + if (!enabled) { + PAIMON_RETURN_NOT_OK(CatalogUtils::CheckPathIsTableLocation( + table_identifier, + location ? Result<std::string>(location.value()) + : catalog->GetTableLocation(table_identifier), + table_path, operation)); + } + } + } + if (!enabled && has_specific_table_schema && + read_branch == BranchManager::DEFAULT_MAIN_BRANCH && !system_table_path) { + SchemaManager schema_manager(file_system, path, read_branch); + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest_schema, + schema_manager.Latest()); + if (latest_schema) { + PAIMON_ASSIGN_OR_RAISE(enabled, SchemaValidation::QueryAuthReadEnabled( + latest_schema.value()->Options(), options)); + } + } + if (!enabled) { + return options; + } + std::map<std::string, std::string> read_options = options; + read_options[Options::QUERY_AUTH_ENABLED] = "true"; + return read_options; +} + /// @param loaded_schema The table's schema, when the caller already read it, or null when it has /// to be read here. `TableScan::Create()` reads one to dispatch on its table type. +/// @param read_options The options of the scan with `ResolveQueryAuthReadOptions()` applied. +/// @param query_auth_applied Whether a `QueryAuthTableScan` applies the rules around this scan, so +/// a table with `query-auth.enabled=true` is planned as it is. Result<std::unique_ptr<TableScan>> NewDataTableScan( - const std::shared_ptr<ScanContext>& context, const std::shared_ptr<TableSchema>& loaded_schema); + const std::shared_ptr<ScanContext>& context, const std::shared_ptr<TableSchema>& loaded_schema, + const std::map<std::string, std::string>& read_options, bool query_auth_applied = false); /// Maps a `ScanContext` onto `FormatTableScan`, which takes one partition filter and a limit and /// nothing else. @@ -720,13 +822,21 @@ if (shared_context->GetRealtimeContext()) { PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateBranch(tmp_options.GetBranch())); } + std::map<std::string, std::string> read_options; + PAIMON_ASSIGN_OR_RAISE( + read_options, + ResolveQueryAuthReadOptions(shared_context->GetCatalog(), shared_context->GetIdentifier(), + shared_context->GetOptions(), + BranchManager::NormalizeBranch(tmp_options.GetBranch()), + tmp_options.GetFileSystem(), shared_context->GetPath(), + shared_context->GetSpecificTableSchema().has_value(), "scan")); PAIMON_ASSIGN_OR_RAISE(std::optional<SystemTablePath> system_table_path, SystemTableLoader::TryParsePath(shared_context->GetPath())); if (system_table_path) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr<SystemTable> system_table, SystemTableLoader::LoadFromPath(tmp_options.GetFileSystem(), shared_context->GetPath(), - shared_context->GetOptions())); + read_options)); return system_table->NewScan(shared_context); } // A format table is planned by listing directories, so it never reaches the manifest path @@ -737,14 +847,13 @@ std::shared_ptr<FormatTable> format_table, FormatTableLoader::TryLoad(tmp_options.GetFileSystem(), shared_context->GetPath(), BranchManager::NormalizeBranch(tmp_options.GetBranch()), - shared_context->GetOptions(), - shared_context->GetSpecificTableSchema(), + read_options, shared_context->GetSpecificTableSchema(), /*schema_manager=*/nullptr, &latest_schema)); if (format_table != nullptr) { return NewFormatTableScan(format_table, shared_context); } // With the schema the dispatch already read, so the managed path does not read it again. - return NewDataTableScan(shared_context, latest_schema); + return NewDataTableScan(shared_context, latest_schema, read_options); } namespace { @@ -792,8 +901,8 @@ } Result<std::unique_ptr<TableScan>> NewDataTableScan( - const std::shared_ptr<ScanContext>& context, - const std::shared_ptr<TableSchema>& loaded_schema) { + const std::shared_ptr<ScanContext>& context, const std::shared_ptr<TableSchema>& loaded_schema, + const std::map<std::string, std::string>& read_options, bool query_auth_applied) { PAIMON_ASSIGN_OR_RAISE( CoreOptions tmp_options, CoreOptions::FromMap(context->GetOptions(), context->GetSpecificFileSystem(), {})); @@ -814,6 +923,16 @@ } table_schema = latest_table_schema.value(); } + PAIMON_ASSIGN_OR_RAISE(bool query_auth, SchemaValidation::QueryAuthReadEnabled( + table_schema->Options(), read_options)); + if (query_auth && !query_auth_applied) { + return QueryAuthTableScan::Create( + context, table_schema, tmp_options.GetFileSystem(), branch, + [table_schema, read_options](const std::shared_ptr<ScanContext>& inner_context) { + return NewDataTableScan(inner_context, table_schema, read_options, + /*query_auth_applied=*/true); + }); + } PAIMON_ASSIGN_OR_RAISE(bool read_optimized, OptionsUtils::GetValueFromMap<bool>(context->GetOptions(), kReadOptimizedScanOption, false)); @@ -912,13 +1031,12 @@ /// @param loaded_schema The schema of `branch`, when the caller already read it, or null when it /// has to be read here: `TableRead::Create()` reads one to dispatch on the table type, and /// reading it again would cost a second listing and read. -Result<std::unique_ptr<InternalReadContext>> CreateInternalReadContext( +Result<std::shared_ptr<TableSchema>> LoadReadSchema( const std::shared_ptr<ReadContext>& context, const std::string& branch, const std::shared_ptr<TableSchema>& loaded_schema) { if (context->GetRealtimeContext()) { PAIMON_RETURN_NOT_OK(RealtimeUtils::ValidateBranch(branch)); } - std::map<std::string, std::string> tmp_options = context->GetOptions(); std::shared_ptr<TableSchema> table_schema; const auto& specific_table_schema = context->GetSpecificTableSchema(); if (loaded_schema != nullptr) { @@ -927,9 +1045,10 @@ PAIMON_ASSIGN_OR_RAISE(table_schema, TableSchema::CreateFromJson(specific_table_schema.value())); } else { - PAIMON_ASSIGN_OR_RAISE(CoreOptions tmp_core_options, - CoreOptions::FromMap(tmp_options, context->GetSpecificFileSystem(), - context->GetFileSystemSchemeToIdentifierMap())); + PAIMON_ASSIGN_OR_RAISE( + CoreOptions tmp_core_options, + CoreOptions::FromMap(context->GetOptions(), context->GetSpecificFileSystem(), + context->GetFileSystemSchemeToIdentifierMap())); SchemaManager schema_manager(tmp_core_options.GetFileSystem(), context->GetPath(), branch); PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> latest_schema, schema_manager.Latest()); @@ -940,10 +1059,16 @@ table_schema = latest_schema.value(); } assert(table_schema); + return table_schema; +} +/// @param table_schema The schema of `branch`, as `LoadReadSchema()` returns it. +Result<std::unique_ptr<InternalReadContext>> CreateInternalReadContext( + const std::shared_ptr<ReadContext>& context, const std::string& branch, + const std::shared_ptr<TableSchema>& table_schema) { // merge options auto options = table_schema->Options(); - for (const auto& [key, value] : tmp_options) { + for (const auto& [key, value] : context->GetOptions()) { options[key] = value; } if (branch != BranchManager::DEFAULT_MAIN_BRANCH) { @@ -975,11 +1100,27 @@ return KeyValueTableRead::Create(path_factory, internal_context, memory_pool, executor); } +/// @param read_options The options of the read with `ResolveQueryAuthReadOptions()` applied. +/// @param query_auth_applied Whether a `QueryAuthTableRead` applies the rules around this read, so +/// a table with `query-auth.enabled=true` is read as it is. Result<std::unique_ptr<TableRead>> NewDataTableRead( - const std::shared_ptr<ReadContext>& context, - const std::shared_ptr<TableSchema>& loaded_schema) { + const std::shared_ptr<ReadContext>& context, const std::shared_ptr<TableSchema>& loaded_schema, + const std::map<std::string, std::string>& read_options, bool query_auth_applied = false) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> table_schema, + LoadReadSchema(context, context->GetBranch(), loaded_schema)); + PAIMON_ASSIGN_OR_RAISE(bool query_auth, SchemaValidation::QueryAuthReadEnabled( + table_schema->Options(), read_options)); + // Dispatched before the read context is built, which consumes the read schema. + if (query_auth && !query_auth_applied) { + return QueryAuthTableRead::Create( + context, table_schema, GetArrowPool(context->GetMemoryPool()), + [table_schema, read_options](std::unique_ptr<ReadContext> inner_context) { + return NewDataTableRead(std::move(inner_context), table_schema, read_options, + /*query_auth_applied=*/true); + }); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InternalReadContext> internal_context, - CreateInternalReadContext(context, context->GetBranch(), loaded_schema)); + CreateInternalReadContext(context, context->GetBranch(), table_schema)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr<TableRead> table_read, @@ -992,10 +1133,34 @@ return std::move(table_read); } + if (query_auth_applied) { + return Status::NotImplemented( + "Fallback reads are not supported on a query-auth table: each branch requires its " + "own authorized splits and schema validation."); + } + // A fallback branch may enable query auth only in the catalog or in a schema newer than the + // cached one, so its setting is resolved as for the branch read. + std::map<std::string, std::string> fallback_read_options; + PAIMON_ASSIGN_OR_RAISE( + fallback_read_options, + ResolveQueryAuthReadOptions( + context->GetCatalog(), context->GetIdentifier(), context->GetOptions(), + scan_fallback_branch.value(), internal_context->GetCoreOptions().GetFileSystem(), + context->GetPath(), context->GetSpecificTableSchema().has_value(), "read")); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> fallback_schema, + LoadReadSchema(context, /*branch=*/scan_fallback_branch.value(), + /*loaded_schema=*/nullptr)); + PAIMON_ASSIGN_OR_RAISE( + bool fallback_query_auth, + SchemaValidation::QueryAuthReadEnabled(fallback_schema->Options(), fallback_read_options)); + if (fallback_query_auth) { + return Status::NotImplemented(fmt::format( + "Cannot read fallback branch '{}', which enables '{}', of a table that does not.", + scan_fallback_branch.value(), Options::QUERY_AUTH_ENABLED)); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr<InternalReadContext> fallback_context, - CreateInternalReadContext(context, /*branch=*/scan_fallback_branch.value(), - /*loaded_schema=*/nullptr)); + CreateInternalReadContext(context, scan_fallback_branch.value(), fallback_schema)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr<TableRead> fallback_table_read, @@ -1045,6 +1210,13 @@ CoreOptions tmp_core_options, CoreOptions::FromMap(context->GetOptions(), context->GetSpecificFileSystem(), context->GetFileSystemSchemeToIdentifierMap())); + std::map<std::string, std::string> read_options; + PAIMON_ASSIGN_OR_RAISE( + read_options, + ResolveQueryAuthReadOptions(context->GetCatalog(), context->GetIdentifier(), + context->GetOptions(), context->GetBranch(), + tmp_core_options.GetFileSystem(), context->GetPath(), + context->GetSpecificTableSchema().has_value(), "read")); PAIMON_ASSIGN_OR_RAISE(std::optional<SystemTablePath> system_table_path, SystemTableLoader::TryParsePath(context->GetPath())); if (system_table_path) { @@ -1052,10 +1224,13 @@ return Status::NotImplemented( "Precise file-index evaluation is not supported for system tables"); } - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr<SystemTable> system_table, - SystemTableLoader::LoadFromPath(tmp_core_options.GetFileSystem(), context->GetPath(), - context->GetOptions())); + // WithBranch() is kept apart from the options, so hand the branch to the system table. + if (!BranchManager::IsMainBranch(context->GetBranch())) { + read_options[Options::BRANCH] = context->GetBranch(); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<SystemTable> system_table, + SystemTableLoader::LoadFromPath(tmp_core_options.GetFileSystem(), + context->GetPath(), read_options)); return system_table->NewRead(context); } @@ -1066,7 +1241,7 @@ PAIMON_ASSIGN_OR_RAISE( std::shared_ptr<FormatTable> format_table, FormatTableLoader::TryLoad(tmp_core_options.GetFileSystem(), context->GetPath(), - context->GetBranch(), context->GetOptions(), + context->GetBranch(), read_options, context->GetSpecificTableSchema(), /*schema_manager=*/nullptr, &latest_schema)); if (format_table != nullptr) { @@ -1074,7 +1249,7 @@ } // With the schema the dispatch already read, so the managed path does not read it again. - return NewDataTableRead(context, latest_schema); + return NewDataTableRead(context, latest_schema, read_options); } namespace {
diff --git a/src/paimon/core/table/table_test.cpp b/src/paimon/core/table/table_test.cpp index 4792a84..60d7aa5 100644 --- a/src/paimon/core/table/table_test.cpp +++ b/src/paimon/core/table/table_test.cpp
@@ -18,29 +18,67 @@ #include "paimon/catalog/table.h" +#include <algorithm> +#include <cstdint> #include <map> #include <memory> #include <optional> #include <string> +#include <utility> #include <vector> #include "arrow/api.h" +#include "arrow/array/concatenate.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/commit_context.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/core/casting/casting_utils.h" #include "paimon/core/catalog/file_system_catalog.h" +#include "paimon/core/catalog/query_auth_catalog.h" +#include "paimon/core/catalog/table_query_auth_result.h" +#include "paimon/core/io/data_file_meta.h" #include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" +#include "paimon/core/snapshot.h" +#include "paimon/core/table/sink/commit_message_impl.h" +#include "paimon/core/table/source/query_auth_split.h" +#include "paimon/core/table/system/system_table.h" +#include "paimon/defs.h" +#include "paimon/executor.h" #include "paimon/file_store_commit.h" #include "paimon/file_store_write.h" +#include "paimon/global_index/bitmap_global_index_result.h" +#include "paimon/global_index/global_index_scan.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" +#include "paimon/reader/batch_reader.h" +#include "paimon/realtime/realtime_context.h" +#include "paimon/record_batch.h" #include "paimon/scan_context.h" #include "paimon/schema/schema.h" +#include "paimon/table/source/plan.h" +#include "paimon/table/source/split.h" #include "paimon/table/source/table_read.h" #include "paimon/table/source/table_scan.h" #include "paimon/testing/mock/mock_catalog.h" +#include "paimon/testing/utils/deletion_vector_test_helper.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/test_helper.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/utils/range.h" #include "paimon/write_context.h" +#include "rapidjson/document.h" +#include "rapidjson/stringbuffer.h" +#include "rapidjson/writer.h" namespace paimon::test { @@ -280,6 +318,8 @@ auto catalog = std::make_shared<MockVersionManagedCatalog>(); catalog->SetFileSystem(dir->GetFileSystem()); + catalog->SetTableSchema(table->LatestSchema()); + catalog->SetTableLocation(dir->Str()); ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> ctx, ScanContextBuilder(dir->Str()).WithCatalog(catalog, identifier).Finish()); @@ -296,6 +336,8 @@ auto catalog = std::make_shared<MockVersionManagedCatalog>(); catalog->SetFileSystem(dir->GetFileSystem()); + catalog->SetTableSchema(table->LatestSchema()); + catalog->SetTableLocation(dir->Str()); ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> ctx, ReadContextBuilder(dir->Str()).WithCatalog(catalog, identifier).Finish()); @@ -303,6 +345,1514 @@ table->NewRead(std::move(ctx))); } +namespace { + +class MockQueryAuthCatalog : public MockVersionManagedCatalog, public QueryAuthCatalog { + public: + Result<std::shared_ptr<TableQueryAuthResult>> AuthTableQuery( + const Identifier& identifier, + const std::optional<std::vector<std::string>>& select) const override { + auth_identifiers.push_back(identifier); + auth_selects.push_back(select); + std::map<std::string, std::string> selected_masking; + for (const auto& [column, transform] : column_masking) { + if (!select || std::find(select->begin(), select->end(), column) != select->end()) { + selected_masking.emplace(column, transform); + } + } + return TableQueryAuthResult::Create(filter, selected_masking); + } + + Result<std::string> GetTableLocation(const Identifier& identifier) const override { + location_identifiers.push_back(identifier); + if (table_location.empty()) { + return MockVersionManagedCatalog::GetTableLocation(identifier); + } + return table_location; + } + + Result<std::shared_ptr<Schema>> LoadTableSchemaAndLocation( + const Identifier& identifier, std::string* location) const override { + *location = table_location; + return LoadTableSchema(identifier); + } + + Result<std::shared_ptr<Schema>> LoadTableSchema(const Identifier& identifier) const override { + Result<std::shared_ptr<Schema>> schema = + MockVersionManagedCatalog::LoadTableSchema(identifier); + auto iter = object_schemas.find(identifier.GetObjectName()); + if (iter != object_schemas.end()) { + return iter->second; + } + return schema; + } + + std::string table_location; + std::map<std::string, std::shared_ptr<Schema>> object_schemas; + std::vector<std::string> filter; + std::map<std::string, std::string> column_masking; + mutable std::vector<Identifier> auth_identifiers; + mutable std::vector<std::optional<std::vector<std::string>>> auth_selects; + mutable std::vector<Identifier> location_identifiers; +}; + +std::shared_ptr<arrow::Schema> QueryAuthTableSchema() { + return arrow::schema({arrow::field("id", arrow::int32(), /*nullable=*/false), + arrow::field("name", arrow::utf8()), arrow::field("phone", arrow::utf8()), + arrow::field("salary", arrow::int64())}); +} + +std::map<std::string, std::string> QueryAuthTableOptions(const std::string& file_format) { + return {{Options::FILE_FORMAT, file_format}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::QUERY_AUTH_ENABLED, "true"}}; +} + +Result<std::vector<std::shared_ptr<CommitMessage>>> WriteQueryAuthRows( + TestHelper* helper, const std::string& rows, int64_t commit_identifier, + const std::vector<RecordBatch::RowKind>& row_kinds = {}) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(QueryAuthTableSchema()->fields()), rows, + /*partition_map=*/{}, /*bucket=*/0, row_kinds)); + return helper->WriteAndCommit(std::move(batch), commit_identifier, std::nullopt); +} + +Result<std::unique_ptr<TestHelper>> CreateQueryAuthTable( + const std::string& root, const std::vector<std::string>& primary_keys, + const std::map<std::string, std::string>& options, const std::vector<std::string>& commits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TestHelper> helper, + TestHelper::Create(root, QueryAuthTableSchema(), /*partition_keys=*/{}, + primary_keys, options, /*is_streaming_mode=*/true)); + for (size_t i = 0; i < commits.size(); ++i) { + PAIMON_RETURN_NOT_OK( + WriteQueryAuthRows(helper.get(), commits[i], static_cast<int64_t>(i)).status()); + } + return helper; +} + +Result<std::unique_ptr<TestHelper>> CreateQueryAuthTable(const std::string& root) { + return CreateQueryAuthTable(root, /*primary_keys=*/{}, QueryAuthTableOptions("orc"), + {R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200], + [3, "carol", null, 300]])"}); +} + +Result<std::shared_ptr<MockQueryAuthCatalog>> CreateQueryAuthCatalog( + const TestHelper& helper, const std::shared_ptr<FileSystem>& file_system) { + PAIMON_ASSIGN_OR_RAISE(std::optional<std::shared_ptr<TableSchema>> table_schema, + helper.LatestSchema()); + if (!table_schema) { + return Status::Invalid("the table has no schema"); + } + auto catalog = std::make_shared<MockQueryAuthCatalog>(); + catalog->table_location = helper.TablePath(); + catalog->SetFileSystem(file_system); + catalog->SetTableSchema(checked_pointer_cast<Schema>(table_schema.value())); + return catalog; +} + +const char kPhonePrefixMask[] = + R"({"name": "SUBSTRING", "inputs": [{"index": 2, "name": "phone", "type": "STRING"}, 1, 3]})"; +const char kSalaryAbove150[] = + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 3,)" + R"( "name": "salary", "type": "BIGINT"}}, "function": "GREATER_THAN", "literals": [150]})"; + +std::string ConstantFilter(const std::string& function) { + return fmt::format(R"({{"kind": "LEAF", "transform": {{"name": "NULL"}},)" + R"( "function": "{}", "literals": []}})", + function); +} + +std::string CompoundFilter(const std::string& function, const std::string& left, + const std::string& right) { + return fmt::format(R"({{"kind": "COMPOUND", "function": "{}", "children": [{}, {}]}})", + function, left, right); +} + +Result<std::shared_ptr<arrow::StructArray>> ReadAll( + std::unique_ptr<ReadContext> context, const std::vector<std::shared_ptr<Split>>& splits) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> read, TableRead::Create(std::move(context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, read->CreateReader(splits)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ChunkedArray> chunks, + ReadResultCollector::CollectResult(std::move(reader))); + read.reset(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> array, + arrow::Concatenate(chunks->chunks(), arrow::default_memory_pool())); + return checked_pointer_cast<arrow::StructArray>(array); +} + +void AssertColumn(const arrow::StructArray& rows, const std::string& name, + const std::string& expected) { + std::shared_ptr<arrow::Array> column = rows.GetFieldByName(name); + ASSERT_NE(nullptr, column) << name; + std::shared_ptr<arrow::Array> expected_column = + arrow::ipc::internal::json::ArrayFromJSON(column->type(), expected).ValueOrDie(); + ASSERT_TRUE(column->Equals(*expected_column)) + << name << ": " << column->ToString() << " != " << expected; +} + +Result<std::shared_ptr<Plan>> PlanQueryAuth(const std::string& path, + const std::shared_ptr<MockQueryAuthCatalog>& catalog, + const std::shared_ptr<Predicate>& predicate, + std::optional<int32_t> limit) { + ScanContextBuilder builder(path); + builder.WithCatalog(catalog, Identifier("foo", "bar")); + if (predicate != nullptr) { + builder.SetPredicate(predicate); + } + if (limit) { + builder.SetLimit(limit.value()); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> context, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> scan, TableScan::Create(std::move(context))); + return scan->CreatePlan(); +} + +Result<std::shared_ptr<arrow::StructArray>> PlanAndReadAll( + const std::string& path, const std::shared_ptr<MockQueryAuthCatalog>& catalog, + const std::vector<std::string>& read_fields = {}) { + ScanContextBuilder scan_builder(path); + scan_builder.WithCatalog(catalog, Identifier("foo", "bar")); + ReadContextBuilder read_builder(path); + if (!read_fields.empty()) { + scan_builder.SetReadFieldNames(read_fields); + read_builder.SetReadFieldNames(read_fields); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> scan_context, scan_builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, scan->CreatePlan()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> read_context, read_builder.Finish()); + return ReadAll(std::move(read_context), plan->Splits()); +} + +} // namespace + +TEST(TableTest, TestQueryAuthTableAppliesTheRulesOfItsCatalog) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + const Identifier identifier("foo", "bar"); + catalog->filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0,)" + R"( "name": "id", "type": "INT NOT NULL"}}, "function": "GREATER_THAN", "literals": [1]})"}; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_TRUE(catalog->auth_identifiers.empty()); + catalog->column_masking = {{"phone", kPhonePrefixMask}, {"salary", R"({"name": "NULL"})"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + ASSERT_EQ(1u, catalog->auth_identifiers.size()); + ASSERT_EQ(identifier, catalog->auth_identifiers[0]); + ASSERT_FALSE(catalog->auth_selects[0]); + ASSERT_FALSE(plan->Splits().empty()); + for (const std::shared_ptr<Split>& split : plan->Splits()) { + ASSERT_NE(nullptr, std::dynamic_pointer_cast<QueryAuthSplit>(split)); + } + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(path).Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + ReadAll(std::move(read_context), plan->Splits())); + AssertColumn(*rows, "id", "[2, 3]"); + AssertColumn(*rows, "name", R"(["bob", "carol"])"); + AssertColumn(*rows, "phone", R"(["139", null])"); + AssertColumn(*rows, "salary", "[null, null]"); + + ASSERT_OK_AND_ASSIGN(read_context, + ReadContextBuilder(path).SetReadFieldNames({"name"}).Finish()); + ASSERT_OK_AND_ASSIGN(rows, ReadAll(std::move(read_context), plan->Splits())); + AssertColumn(*rows, "name", R"(["bob", "carol"])"); + ASSERT_EQ(nullptr, rows->GetFieldByName("id")); + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path).SetReadFieldIds({1}).Finish()); + ASSERT_OK_AND_ASSIGN(rows, ReadAll(std::move(read_context), plan->Splits())); + AssertColumn(*rows, "name", R"(["bob", "carol"])"); + ASSERT_EQ(nullptr, rows->GetFieldByName("id")); + + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path) + .SetReadFieldNames({"name", "phone"}) + .SetPredicate(PredicateBuilder::Equal( + /*field_index=*/1, "phone", FieldType::STRING, + Literal(FieldType::STRING, "139", 3))) + .EnablePredicateFilter(true) + .Finish()); + ASSERT_OK_AND_ASSIGN(rows, ReadAll(std::move(read_context), plan->Splits())); + AssertColumn(*rows, "name", R"(["bob"])"); + AssertColumn(*rows, "phone", R"(["139"])"); + ASSERT_EQ(nullptr, rows->GetFieldByName("id")); + + std::vector<std::shared_ptr<Split>> deserialized; + for (const std::shared_ptr<Split>& split : plan->Splits()) { + ASSERT_OK_AND_ASSIGN(std::string bytes, Split::Serialize(split, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> copy, + Split::Deserialize(bytes.data(), bytes.size(), GetDefaultPool())); + auto auth_split = std::dynamic_pointer_cast<QueryAuthSplit>(copy); + ASSERT_NE(nullptr, auth_split); + const auto& rules = std::dynamic_pointer_cast<QueryAuthSplit>(split)->GetAuthResult(); + ASSERT_TRUE(auth_split->GetAuthResult()->SameFilter(*rules)); + ASSERT_EQ(rules->ColumnMasking(), auth_split->GetAuthResult()->ColumnMasking()); + const auto& original = std::dynamic_pointer_cast<QueryAuthSplit>(split); + ASSERT_TRUE(auth_split->GetScope()); + ASSERT_EQ(original->GetScope(), auth_split->GetScope()); + deserialized.push_back(copy); + } + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path).Finish()); + ASSERT_OK_AND_ASSIGN(rows, ReadAll(std::move(read_context), deserialized)); + AssertColumn(*rows, "phone", R"(["139", null])"); + + const auto& planned = std::dynamic_pointer_cast<QueryAuthSplit>(plan->Splits()[0]); + ASSERT_OK_AND_ASSIGN( + std::string java_planned, + Split::Serialize(std::make_shared<QueryAuthSplit>(planned->GetSplit(), + planned->GetAuthResult(), std::nullopt), + GetDefaultPool())); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<Split> java_planned_split, + Split::Deserialize(java_planned.data(), java_planned.size(), GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(read_context, + ReadContextBuilder(path).SetReadFieldNames({"name"}).Finish()); + ASSERT_NOK_WITH_MSG(ReadAll(std::move(read_context), {java_planned_split}), + "The split does not record which columns its query auth rules were " + "authorized for"); + + catalog->filter = {}; + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), "Query auth rules changed"); + ASSERT_EQ(2u, catalog->auth_identifiers.size()); + + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(catalog, identifier) + .SetPartitionFilter({{{"phone", "139"}}}) + .Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), "masked partition key 'phone'"); + + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG( + read->CreateReader( + std::dynamic_pointer_cast<QueryAuthSplit>(plan->Splits()[0])->GetSplit()), + "must be planned by a scan of the table through its catalog"); + ASSERT_NOK(read->CreateCountReader(plan->Splits())); + ASSERT_EQ((std::map<std::string, int32_t>{{"id", 0}, {"phone", 2}, {"salary", 3}}), + planned->GetScope()->rule_field_ids); + ASSERT_EQ((std::map<std::string, int32_t>{{"id", 0}, {"name", 1}, {"phone", 2}, {"salary", 3}}), + planned->GetScope()->authorized_field_ids); + QueryAuthScope replaced_scope = planned->GetScope().value(); + replaced_scope.rule_field_ids["phone"] = 1; + ASSERT_NOK_WITH_MSG( + read->CreateReader(std::make_shared<QueryAuthSplit>( + planned->GetSplit(), planned->GetAuthResult(), replaced_scope)), + "(schema 0): Query auth rules reference column 'phone' which the schema being read " + "exposes as a different column of the same name"); + QueryAuthScope unbound_scope = planned->GetScope().value(); + unbound_scope.rule_field_ids.erase("phone"); + ASSERT_NOK_WITH_MSG(read->CreateReader(std::make_shared<QueryAuthSplit>( + planned->GetSplit(), planned->GetAuthResult(), unbound_scope)), + "Query auth rules reference column 'phone', which the schema the scan " + "planned with did not hold"); + + catalog->filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 9,)" + R"( "name": "gone", "type": "INT"}}, "function": "IS_NULL", "literals": []})"}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), + "Query auth of Identifier{database='foo', object='bar'} (branch main, " + "schema 0): Row filter references column 'gone' which does not exist"); + catalog->filter = {}; + catalog->column_masking = {{"name", R"({"name": "UPPER", "inputs": []})"}}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), "Column mask transform 'UPPER' is not supported"); + catalog->column_masking = {}; + for (const std::string literals : {"", R"(, "literals": null)"}) { + catalog->filter = {fmt::format( + R"({{"kind": "LEAF", "transform": {{"name": "FIELD_REF", "fieldRef": {{"index": 0,)" + R"( "name": "id", "type": "INT NOT NULL"}}}}, "function": "NOT_IN"{}}})", + literals)}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), + "Query auth row filter NOT_IN requires a non-null 'literals' array."); + } +} + +TEST(TableTest, TestQueryAuthTableRefusesColumnsChangedSinceThePlan) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, + helper->LatestSchema()); + ASSERT_TRUE(table_schema); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + const Identifier identifier("foo", "bar"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->column_masking = {{"phone", R"({"name": "NULL"})"}}; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path) + .WithCatalog(catalog, identifier) + .SetReadFieldNames({"name"}) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> name_plan, scan->CreatePlan()); + ASSERT_FALSE(std::dynamic_pointer_cast<QueryAuthSplit>(name_plan->Splits()[0]) + ->GetAuthResult() + ->HasRules()); + catalog->column_masking = {}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> all_plan, scan->CreatePlan()); + catalog->column_masking = {{"phone", R"({"name": "NULL"})"}}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> masked_plan, scan->CreatePlan()); + catalog->column_masking = {}; + catalog->filter = {CompoundFilter( + "OR", ConstantFilter("TRUE"), + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 2,)" + R"( "name": "phone", "type": "STRING"}}, "function": "IS_NULL", "literals": []})")}; + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(scan, TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> folded_plan, scan->CreatePlan()); + catalog->filter = {}; + catalog->column_masking = {{"phone", R"({"name": "NULL"})"}}; + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(catalog, identifier) + .SetReadFieldNames({"name"}) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> unplanned_scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(catalog, identifier) + .SetReadFieldNames({"id"}) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> id_scan, + TableScan::Create(std::move(scan_context))); + + ASSERT_OK_AND_ASSIGN(std::string schema_json, table_schema.value()->ToJsonString()); + rapidjson::Document document; + document.Parse(schema_json.c_str()); + ASSERT_FALSE(document.HasParseError()); + auto& allocator = document.GetAllocator(); + document["id"].SetInt64(1); + rapidjson::Value& fields = document["fields"]; + fields[1]["name"].SetString("nickname", allocator); + fields[2]["name"].SetString("name", allocator); + rapidjson::Value secret(rapidjson::kObjectType); + secret.AddMember("id", 4, allocator); + secret.AddMember("name", "secret", allocator); + secret.AddMember("type", "STRING", allocator); + fields.PushBack(secret, allocator); + document["highestFieldId"].SetInt(4); + rapidjson::StringBuffer buffer; + rapidjson::Writer<rapidjson::StringBuffer> writer(buffer); + document.Accept(writer); + ASSERT_OK(TableSchema::CreateFromJson(buffer.GetString()).status()); + ASSERT_OK(dir->GetFileSystem()->AtomicStore( + PathUtil::JoinPath(SchemaManager(dir->GetFileSystem(), path).SchemaDirectory(), "schema-1"), + buffer.GetString())); + + ASSERT_NOK_WITH_MSG(unplanned_scan->CreatePlan(), + "Column 'name' is field 1 in schema 0 the scan plans with, but field 2 in " + "the latest schema 1"); + catalog->filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 4,)" + R"( "name": "secret", "type": "STRING"}}, "function": "IS_NULL", "literals": []})"}; + ASSERT_NOK_WITH_MSG(id_scan->CreatePlan(), + "Query auth rules reference column 'secret', which the schema the scan " + "planned with did not hold"); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(path).SetReadFieldNames({"name"}).Finish()); + ASSERT_NOK_WITH_MSG(ReadAll(std::move(read_context), name_plan->Splits()), + "Column 'name' is read as field 2, but the scan authorized field 1 under " + "that name"); + ASSERT_OK_AND_ASSIGN(read_context, + ReadContextBuilder(path).SetReadFieldNames({"secret"}).Finish()); + ASSERT_NOK_WITH_MSG(ReadAll(std::move(read_context), all_plan->Splits()), + "Column 'secret' is read, but the scan did not ask the catalog to " + "authorize it"); + ASSERT_OK_AND_ASSIGN(read_context, + ReadContextBuilder(path).SetReadFieldNames({"id", "salary"}).Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + ReadAll(std::move(read_context), all_plan->Splits())); + AssertColumn(*rows, "id", "[1, 2, 3]"); + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path).SetReadFieldNames({"id"}).Finish()); + ASSERT_NOK_WITH_MSG(ReadAll(std::move(read_context), masked_plan->Splits()), + "Query auth rules reference column 'phone' which the schema being read " + "exposes as 'name' (renamed since)"); + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path).SetReadFieldNames({"id"}).Finish()); + ASSERT_NOK_WITH_MSG(ReadAll(std::move(read_context), folded_plan->Splits()), + "Query auth rules reference column 'phone' which the schema being read " + "exposes as 'name' (renamed since)"); +} + +TEST(TableTest, TestQueryAuthTableAsksForTheRulesOfTheBranchPlanned) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + const std::map<std::string, std::string> table_options = QueryAuthTableOptions("orc"); + std::map<std::string, std::string> dev_options = table_options; + dev_options[Options::BRANCH] = "dev"; + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, "dev") + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, /*primary_keys=*/{}, + table_options) + .status()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> dev_helper, + TestHelper::Create(path, dev_options, /*is_streaming_mode=*/false)); + ASSERT_OK(WriteQueryAuthRows(dev_helper.get(), R"([[7, "dave", "13700007777", 700]])", + /*commit_identifier=*/0) + .status()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->column_masking = {{"phone", R"({"name": "NULL"})"}}; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path) + .WithCatalog(catalog, Identifier("foo", "bar")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ(Identifier("foo", "bar$branch_dev"), scan_context->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + ASSERT_FALSE(catalog->TableFileSystemRequests().empty()); + for (const Identifier& requested : catalog->TableFileSystemRequests()) { + ASSERT_EQ(Identifier("foo", "bar$branch_dev"), requested); + } + ASSERT_EQ(std::vector<Identifier>({Identifier("foo", "bar$branch_dev")}), + catalog->auth_identifiers); + ASSERT_EQ(std::vector<Identifier>({Identifier("foo", "bar$branch_dev")}), + catalog->location_identifiers); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(path).AddOption(Options::BRANCH, "dev").Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + ReadAll(std::move(read_context), plan->Splits())); + AssertColumn(*rows, "id", "[7]"); + AssertColumn(*rows, "phone", "[null]"); +} + +TEST(TableTest, TestQueryAuthTableIsNotSkippedThroughAStaleCachedSchema) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, + helper->LatestSchema()); + ASSERT_TRUE(table_schema); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::string schema_json, table_schema.value()->ToJsonString()); + rapidjson::Document document; + document.Parse(schema_json.c_str()); + ASSERT_FALSE(document.HasParseError()); + document["options"].RemoveMember(Options::QUERY_AUTH_ENABLED); + rapidjson::StringBuffer buffer; + rapidjson::Writer<rapidjson::StringBuffer> writer(buffer); + document.Accept(writer); + const std::string stale_schema = buffer.GetString(); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path).SetTableSchema(stale_schema).Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "without the catalog that serves its row filters and column masks"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + ASSERT_OK_AND_ASSIGN( + scan_context, + ScanContextBuilder(path).WithCatalog(catalog, Identifier("foo", "bar")).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(path).SetTableSchema(stale_schema).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(read_context))); + ASSERT_NOK_WITH_MSG( + read->CreateReader( + std::dynamic_pointer_cast<QueryAuthSplit>(plan->Splits()[0])->GetSplit()), + "must be planned by a scan of the table through its catalog"); + + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, "dev") + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::QUERY_AUTH_ENABLED, "false"}}) + .status()); + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(path) + .WithBranch("dev") + .SetTableSchema(stale_schema) + .AddOption(Options::SCAN_FALLBACK_BRANCH, "main") + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), + "Cannot read fallback branch 'main', which enables 'query-auth.enabled'"); +} + +TEST(TableTest, TestQueryAuthScanDoesNotPruneByTheRawValuesOfAMaskedColumn) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::map<std::string, std::string> options = QueryAuthTableOptions("orc"); + options[Options::SOURCE_SPLIT_TARGET_SIZE] = "1 kb"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{}, options, + {R"([[1, "alice", "13800001111", 100]])", + R"([[3, "carol", "13700003333", 300]])", + R"([[2, "bob", "13900002222", 200]])"})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + auto read_with = + [&](const std::shared_ptr<Predicate>& predicate, const std::shared_ptr<Plan>& plan, + bool enable_predicate_filter = true) -> Result<std::shared_ptr<arrow::StructArray>> { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> context, + ReadContextBuilder(path) + .SetPredicate(predicate) + .EnablePredicateFilter(enable_predicate_filter) + .Finish()); + return ReadAll(std::move(context), plan->Splits()); + }; + const std::shared_ptr<Predicate> phone_is_139 = PredicateBuilder::Equal( + /*field_index=*/2, "phone", FieldType::STRING, Literal(FieldType::STRING, "139", 3)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path) + .WithCatalog(catalog, Identifier("foo", "bar")) + .SetPredicate(phone_is_139) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + ASSERT_TRUE(plan->Splits().empty()); + + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), "Query auth rules changed"); + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, phone_is_139, std::nullopt)); + ASSERT_EQ(3u, plan->Splits().size()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, read_with(phone_is_139, plan)); + AssertColumn(*rows, "id", "[2]"); + AssertColumn(*rows, "phone", R"(["139"])"); + + const std::shared_ptr<Predicate> id_is_1 = PredicateBuilder::Equal( + /*field_index=*/0, "id", FieldType::INT, Literal(static_cast<int32_t>(1))); + ASSERT_OK_AND_ASSIGN(rows, read_with(phone_is_139, plan, /*enable_predicate_filter=*/false)); + AssertColumn(*rows, "id", "[2]"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> phone_and_id, + PredicateBuilder::And({phone_is_139, id_is_1})); + ASSERT_OK_AND_ASSIGN(rows, read_with(phone_and_id, plan, /*enable_predicate_filter=*/false)); + AssertColumn(*rows, "id", "[2]"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Predicate> phone_or_id, + PredicateBuilder::Or({phone_is_139, id_is_1})); + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, phone_or_id, std::nullopt)); + ASSERT_EQ(3u, plan->Splits().size()); + ASSERT_OK_AND_ASSIGN(rows, read_with(phone_or_id, plan)); + AssertColumn(*rows, "id", "[1, 2]"); + AssertColumn(*rows, "phone", R"(["138", "139"])"); + + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, phone_is_139, /*limit=*/1)); + ASSERT_EQ(3u, plan->Splits().size()); + ASSERT_OK_AND_ASSIGN(rows, read_with(phone_is_139, plan)); + AssertColumn(*rows, "id", "[2]"); + AssertColumn(*rows, "phone", R"(["139"])"); + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, /*predicate=*/nullptr, /*limit=*/1)); + ASSERT_EQ(1u, plan->Splits().size()); + catalog->filter = {kSalaryAbove150}; + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, /*predicate=*/nullptr, /*limit=*/1)); + ASSERT_EQ(2u, plan->Splits().size()); +} + +TEST(TableTest, TestQueryAuthReadRefusesAColumnTheScanDidNotAuthorize) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + auto read_names = [&](bool enable_predicate_filter, + const std::vector<std::shared_ptr<Split>>& splits) + -> Result<std::shared_ptr<arrow::StructArray>> { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<ReadContext> context, + ReadContextBuilder(path) + .SetReadFieldNames({"name"}) + .SetPredicate(PredicateBuilder::Equal(/*field_index=*/1, "phone", FieldType::STRING, + Literal(FieldType::STRING, "139", 3))) + .EnablePredicateFilter(enable_predicate_filter) + .Finish()); + return ReadAll(std::move(context), splits); + }; + const std::string unauthorized = + "Column 'phone' is read, but the scan did not ask the catalog to authorize it"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> name_scan_context, + ScanContextBuilder(path) + .WithCatalog(catalog, Identifier("foo", "bar")) + .SetReadFieldNames({"name"}) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> name_scan, + TableScan::Create(std::move(name_scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> name_plan, name_scan->CreatePlan()); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{"name"}), + catalog->auth_selects.back()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> all_columns_context, + ReadContextBuilder(path).Finish()); + ASSERT_NOK_WITH_MSG( + ReadAll(std::move(all_columns_context), name_plan->Splits()), + "Column 'id' is read, but the scan did not ask the catalog to authorize it"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> name_context, + ReadContextBuilder(path).SetReadFieldNames({"name"}).Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + ReadAll(std::move(name_context), name_plan->Splits())); + AssertColumn(*rows, "name", R"(["alice", "bob", "carol"])"); + ASSERT_NOK_WITH_MSG(read_names(/*enable_predicate_filter=*/true, name_plan->Splits()), + unauthorized); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, + PlanQueryAuth(path, catalog, /*predicate=*/nullptr, std::nullopt)); + ASSERT_OK_AND_ASSIGN(rows, read_names(/*enable_predicate_filter=*/false, plan->Splits())); + AssertColumn(*rows, "name", R"(["bob"])"); + ASSERT_EQ(nullptr, rows->GetFieldByName("phone")); + auto planned = std::dynamic_pointer_cast<QueryAuthSplit>(plan->Splits()[0]); + ASSERT_NE(nullptr, planned); + ASSERT_EQ(1u, planned->GetAuthResult()->MaskedFields().count("phone")); + QueryAuthScope scope = planned->GetScope().value(); + scope.authorized_field_ids.erase("phone"); + ASSERT_NOK_WITH_MSG(read_names(/*enable_predicate_filter=*/false, + {std::make_shared<QueryAuthSplit>( + planned->GetSplit(), planned->GetAuthResult(), scope)}), + unauthorized); +} + +TEST(TableTest, TestQueryAuthScanPushesTheLimitPastARowFilterOnPartitionKeys) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> helper, + TestHelper::Create(dir->Str(), QueryAuthTableSchema(), /*partition_keys=*/{"name"}, + /*primary_keys=*/{}, QueryAuthTableOptions("orc"), + /*is_streaming_mode=*/true)); + const std::vector<std::pair<std::string, std::string>> partitions = { + {"alice", R"([[1, "alice", "13800001111", 100]])"}, + {"bob", R"([[2, "bob", "13900002222", 200]])"}, + {"carol", R"([[3, "carol", "13700003333", 300]])"}}; + for (size_t i = 0; i < partitions.size(); ++i) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(QueryAuthTableSchema()->fields()), + partitions[i].second, {{"name", partitions[i].first}}, + /*bucket=*/0, /*row_kinds=*/{})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), static_cast<int64_t>(i), std::nullopt) + .status()); + } + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 1,)" + R"( "name": "name", "type": "STRING"}}, "function": "IN", "literals": ["alice", "bob"]})"}; + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, + PlanQueryAuth(path, catalog, /*predicate=*/nullptr, std::nullopt)); + ASSERT_EQ(2u, plan->Splits().size()); + ASSERT_OK_AND_ASSIGN(plan, PlanQueryAuth(path, catalog, /*predicate=*/nullptr, /*limit=*/1)); + ASSERT_EQ(1u, plan->Splits().size()); +} + +TEST(TableTest, TestQueryAuthTableAppliesTheSameRulesToEachFileFormat) { + for (const std::string file_format : {"orc", "parquet"}) { + SCOPED_TRACE(file_format); + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{}, + QueryAuthTableOptions(file_format), + {R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200], + [3, "carol", null, 300]])"})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->filter = {kSalaryAbove150}; + catalog->column_masking = {{"phone", kPhonePrefixMask}, {"salary", R"({"name": "NULL"})"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + PlanAndReadAll(PathUtil::JoinPath(dir->Str(), "foo.db/bar"), catalog)); + AssertColumn(*rows, "id", "[2, 3]"); + AssertColumn(*rows, "name", R"(["bob", "carol"])"); + AssertColumn(*rows, "phone", R"(["139", null])"); + AssertColumn(*rows, "salary", "[null, null]"); + } +} + +TEST(TableTest, TestQueryAuthRefusesToIgnoreAMaskOnALevelColumn) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto schema = arrow::schema({arrow::field("id", arrow::int32(), /*nullable=*/false), + arrow::field("_LEVEL", arrow::utf8())}); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> helper, + TestHelper::Create(dir->Str(), schema, /*partition_keys=*/{}, /*primary_keys=*/{}, + QueryAuthTableOptions("orc"), /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(schema->fields()), R"([[1, "secret"]])", + /*partition_map=*/{}, /*bucket=*/0, /*row_kinds=*/{})); + ASSERT_OK( + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, std::nullopt).status()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, PlanAndReadAll(path, catalog)); + AssertColumn(*rows, "_LEVEL", R"(["secret"])"); + catalog->column_masking = {{"_LEVEL", R"({"name": "NULL"})"}}; + ASSERT_NOK_WITH_MSG(PlanAndReadAll(path, catalog), "refusing to ignore its mask"); +} + +TEST(TableTest, TestQueryAuthTableReadsNoRowsUnderAJavaAlwaysFalseFilter) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + auto read_ids = + [&](const std::vector<std::string>& read_fields) -> Result<std::shared_ptr<arrow::Array>> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, + PlanQueryAuth(path, catalog, /*predicate=*/nullptr, std::nullopt)); + ReadContextBuilder builder(path); + builder.SetReadFieldNames(read_fields); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> context, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, + read->CreateReader(plan->Splits())); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::ChunkedArray> chunks, + ReadResultCollector::CollectResult(std::move(reader))); + if (chunks == nullptr) { + return std::shared_ptr<arrow::Array>(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> array, + arrow::Concatenate(chunks->chunks(), arrow::default_memory_pool())); + return checked_pointer_cast<arrow::StructArray>(array)->GetFieldByName("id"); + }; + for (const auto& [filter, expected] : + std::vector<std::pair<std::vector<std::string>, std::optional<std::string>>>{ + {{ConstantFilter("TRUE")}, "[1, 2, 3]"}, + {{CompoundFilter("OR", ConstantFilter("FALSE"), kSalaryAbove150)}, "[2, 3]"}, + {{ConstantFilter("FALSE")}, std::nullopt}, + {{kSalaryAbove150, + CompoundFilter("AND", ConstantFilter("TRUE"), ConstantFilter("FALSE"))}, + std::nullopt}}) { + SCOPED_TRACE(filter[0]); + catalog->filter = filter; + for (const std::vector<std::string>& read_fields : + {std::vector<std::string>{"id"}, std::vector<std::string>{"id", "name"}}) { + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> ids, read_ids(read_fields)); + if (!expected) { + ASSERT_EQ(nullptr, ids); + continue; + } + ASSERT_NE(nullptr, ids); + ASSERT_TRUE(ids->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), expected.value()) + .ValueOrDie())) + << ids->ToString(); + } + } +} + +TEST(TableTest, TestQueryAuthPrimaryKeyTableAuthorizesTheMergedRows) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{"id"}, QueryAuthTableOptions("orc"), + {R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200], + [3, "carol", "13700003333", 300]])", + R"([[2, "bob", "13900002222", 50]])"})); + ASSERT_OK(WriteQueryAuthRows(helper.get(), R"([[3, "carol", "13700003333", 50]])", + /*commit_identifier=*/2, {RecordBatch::RowKind::DELETE}) + .status()); + ASSERT_OK(WriteQueryAuthRows(helper.get(), R"([[1, "alice", "13800001111", 400]])", + /*commit_identifier=*/3) + .status()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->filter = {kSalaryAbove150}; + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + PlanAndReadAll(PathUtil::JoinPath(dir->Str(), "foo.db/bar"), catalog)); + AssertColumn(*rows, "id", "[1]"); + AssertColumn(*rows, "phone", R"(["138"])"); + AssertColumn(*rows, "salary", "[400]"); +} + +TEST(TableTest, TestQueryAuthDoesNotBringBackRowsADeletionVectorDeletes) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{}, + {{Options::FILE_FORMAT, "orc"}, + {Options::BUCKET, "-1"}, + {Options::DELETION_VECTORS_ENABLED, "true"}, + {Options::QUERY_AUTH_ENABLED, "true"}}, + /*commits=*/{})); + ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> data_messages, + WriteQueryAuthRows(helper.get(), R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200], + [3, "carol", "13700003333", 300]])", + /*commit_identifier=*/0)); + ASSERT_EQ(1u, data_messages.size()); + auto data_message = std::dynamic_pointer_cast<CommitMessageImpl>(data_messages[0]); + ASSERT_NE(nullptr, data_message); + ASSERT_EQ(1u, data_message->GetNewFilesIncrement().NewFiles().size()); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<CommitMessage> deletion_message, + DeletionVectorTestHelper::CreateDeletionVectorCommitMessage( + dir->GetFileSystem(), path, "orc", data_messages[0], + {{data_message->GetNewFilesIncrement().NewFiles()[0]->file_name, {1}}}, + GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommitContext> commit_context, + CommitContextBuilder(path, "commit_user").Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreCommit> commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK(commit->Commit({deletion_message}, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->filter = {kSalaryAbove150}; + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, PlanAndReadAll(path, catalog)); + AssertColumn(*rows, "id", "[3]"); + AssertColumn(*rows, "phone", R"(["137"])"); +} + +TEST(TableTest, TestQueryAuthAppliesTheRulesToTheColumnsADataEvolutionTableRebuilds) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{}, + {{Options::FILE_FORMAT, "orc"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}, + {Options::QUERY_AUTH_ENABLED, "true"}}, + {R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200]])"})); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<WriteContext> write_context, + WriteContextBuilder(path, "commit_user").WithWriteSchema({"phone", "salary"}).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreWrite> write, + FileStoreWrite::Create(std::move(write_context))); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_({arrow::field("phone", arrow::utf8()), + arrow::field("salary", arrow::int64())}), + R"([["13700001111", 300], ["13600002222", 50]])", + /*partition_map=*/{}, /*bucket=*/0, + /*row_kinds=*/{})); + ASSERT_OK(write->Write(std::move(batch))); + ASSERT_OK_AND_ASSIGN(std::vector<std::shared_ptr<CommitMessage>> messages, + write->PrepareCommit(/*wait_compaction=*/false, /*commit_identifier=*/1)); + ASSERT_OK(write->Close()); + for (const std::shared_ptr<CommitMessage>& message : messages) { + auto message_impl = std::dynamic_pointer_cast<CommitMessageImpl>(message); + ASSERT_NE(nullptr, message_impl); + for (const std::shared_ptr<DataFileMeta>& file : + message_impl->GetNewFilesIncrement().NewFiles()) { + file->AssignFirstRowId(0); + } + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommitContext> commit_context, + CommitContextBuilder(path, "commit_user").Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileStoreCommit> commit, + FileStoreCommit::Create(std::move(commit_context))); + ASSERT_OK(commit->Commit(messages, /*commit_identifier=*/1)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + catalog->filter = {kSalaryAbove150}; + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + PlanAndReadAll(path, catalog, {"name", "phone"})); + AssertColumn(*rows, "name", R"(["alice"])"); + AssertColumn(*rows, "phone", R"(["137"])"); + ASSERT_EQ(nullptr, rows->GetFieldByName("salary")); + ASSERT_EQ(nullptr, rows->GetFieldByName("id")); +} + +TEST(TableTest, TestQueryAuthFiltersDictionaryEncodedParquetColumns) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::map<std::string, std::string> options = QueryAuthTableOptions("parquet"); + options["parquet.read.enable-dictionary-passthrough"] = "true"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + CreateQueryAuthTable(dir->Str(), /*primary_keys=*/{}, options, + {R"([[1, "alice", "13800001111", 100], + [2, "bob", "13900002222", 200], + [3, "carol", null, 300], + [4, "dave", "13700004444", 400]])"})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + auto read = [&](const std::shared_ptr<Predicate>& predicate, + bool enable_predicate_filter) -> Result<std::shared_ptr<arrow::StructArray>> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<Plan> plan, + PlanQueryAuth(path, catalog, predicate, std::nullopt)); + ReadContextBuilder builder(path); + if (predicate != nullptr) { + builder.SetPredicate(predicate).EnablePredicateFilter(enable_predicate_filter); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReadContext> context, builder.Finish()); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<TableRead> table_read, + TableRead::Create(std::move(context))); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader, + table_read->CreateReader(plan->Splits())); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return Status::Invalid("no rows read"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr<arrow::Array> array, + arrow::ImportArray(batch.first.get(), batch.second.get())); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch eof, reader->NextBatch()); + if (!BatchReader::IsEofBatch(eof)) { + ReaderUtils::ReleaseReadBatch(std::move(eof)); + return Status::Invalid("more than one batch read"); + } + reader->Close(); + return checked_pointer_cast<arrow::StructArray>(array); + }; + auto assert_column = [](const arrow::StructArray& rows, const std::string& name, + arrow::Type::type type, const std::string& values) { + std::shared_ptr<arrow::Array> column = rows.GetFieldByName(name); + ASSERT_NE(nullptr, column) << name; + ASSERT_EQ(type, column->type_id()) << name << ": " << column->type()->ToString(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> decoded, + CastingUtils::DecodeDictionary(column, arrow::default_memory_pool())); + ASSERT_TRUE(decoded->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(decoded->type(), values).ValueOrDie())) + << name << ": " << decoded->ToString(); + }; + auto equal = [](int32_t index, const std::string& name, const std::string& value) { + return PredicateBuilder::Equal(index, name, FieldType::STRING, + Literal(FieldType::STRING, value.data(), value.size())); + }; + + catalog->filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 1,)" + R"( "name": "name", "type": "STRING"}}, "function": "NOT_EQUAL", "literals": ["bob"]})"}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + read(/*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + assert_column(*rows, "id", arrow::Type::INT32, "[1, 3, 4]"); + assert_column(*rows, "name", arrow::Type::STRING, R"(["alice", "carol", "dave"])"); + // Decoded: the dictionary of the column would still hold bob's phone number. + assert_column(*rows, "phone", arrow::Type::STRING, R"(["13800001111", null, "13700004444"])"); + + catalog->filter = {kSalaryAbove150}; + ASSERT_OK_AND_ASSIGN(rows, read(equal(1, "name", "carol"), /*enable_predicate_filter=*/true)); + assert_column(*rows, "id", arrow::Type::INT32, "[3]"); + assert_column(*rows, "name", arrow::Type::STRING, R"(["carol"])"); + assert_column(*rows, "phone", arrow::Type::STRING, "[null]"); + + catalog->column_masking = {{"phone", kPhonePrefixMask}}; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<Predicate> phone_or_name, + PredicateBuilder::Or({equal(2, "phone", "137"), equal(1, "name", "carol")})); + for (bool enable_predicate_filter : {false, true}) { + SCOPED_TRACE(enable_predicate_filter); + ASSERT_OK_AND_ASSIGN(rows, read(phone_or_name, enable_predicate_filter)); + assert_column(*rows, "id", arrow::Type::INT32, "[3, 4]"); + assert_column(*rows, "name", arrow::Type::STRING, R"(["carol", "dave"])"); + assert_column(*rows, "phone", arrow::Type::STRING, R"([null, "137"])"); + } + + catalog->filter = {}; + catalog->column_masking = {}; + ASSERT_OK_AND_ASSIGN(rows, read(/*predicate=*/nullptr, /*enable_predicate_filter=*/false)); + assert_column(*rows, "phone", arrow::Type::DICTIONARY, + R"(["13800001111", "13900002222", null, "13700004444"])"); +} + +TEST(TableTest, TestQueryAuthScanReadsOnlyTheLocationOfTheAuthorizedTable) { + auto authorized_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(authorized_dir); + auto other_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(other_dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> authorized, + CreateQueryAuthTable(authorized_dir->Str())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> other, CreateQueryAuthTable(other_dir->Str())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*authorized, authorized_dir->GetFileSystem())); + const Identifier identifier("foo", "bar"); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(other->TablePath()).WithCatalog(catalog, identifier).Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + fmt::format("Invalid: The scan targets {}, but the catalog keeps table " + "foo.bar, whose query auth settings it would apply, at {}", + other->TablePath(), authorized->TablePath())); + ASSERT_TRUE(catalog->auth_identifiers.empty()); + + for (const std::string& path : + {authorized->TablePath() + "/", + PathUtil::JoinPath(authorized_dir->Str() + "/", "foo.db//bar")}) { + SCOPED_TRACE(path); + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK(scan->CreatePlan().status()); + } + + // The identifier is rebound to another table between the scan and its first plan. + ASSERT_OK_AND_ASSIGN( + scan_context, + ScanContextBuilder(authorized->TablePath()).WithCatalog(catalog, identifier).Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> rebound_scan, + TableScan::Create(std::move(scan_context))); + catalog->table_location = other->TablePath(); + const size_t auth_requests = catalog->auth_identifiers.size(); + ASSERT_NOK_WITH_MSG(rebound_scan->CreatePlan(), + fmt::format("The scan targets {}, but the catalog keeps table foo.bar, " + "whose query auth settings it would apply, at {}", + authorized->TablePath(), other->TablePath())); + ASSERT_EQ(auth_requests + 1, catalog->auth_identifiers.size()); + + catalog->table_location.clear(); + ASSERT_OK_AND_ASSIGN( + scan_context, + ScanContextBuilder(authorized->TablePath()).WithCatalog(catalog, identifier).Finish()); + ASSERT_NOK_WITH_MSG( + TableScan::Create(std::move(scan_context)), + "Cannot check that " + authorized->TablePath() + " is the location of table foo.bar"); +} + +TEST(TableTest, TestQueryAuthReadRejectsSplitsAuthorizedForAnotherTable) { + auto authorized_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(authorized_dir); + auto other_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(other_dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> authorized, + CreateQueryAuthTable(authorized_dir->Str())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> other, CreateQueryAuthTable(other_dir->Str())); + // Restored or copied tables can have the same field ids and data file names. + ASSERT_TRUE(TestUtil::CopyDirectory(authorized->TablePath(), other->TablePath())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> authorized_catalog, + CreateQueryAuthCatalog(*authorized, authorized_dir->GetFileSystem())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> other_catalog, + CreateQueryAuthCatalog(*other, other_dir->GetFileSystem())); + other_catalog->column_masking = {{"phone", R"({"name": "NULL"})"}}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> masked, + PlanAndReadAll(other->TablePath(), other_catalog)); + AssertColumn(*masked, "phone", "[null, null, null]"); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> authorized_plan, + PlanQueryAuth(authorized->TablePath(), authorized_catalog, + /*predicate=*/nullptr, std::nullopt)); + std::vector<std::shared_ptr<Split>> splits; + for (const std::shared_ptr<Split>& split : authorized_plan->Splits()) { + ASSERT_OK_AND_ASSIGN(std::string bytes, Split::Serialize(split, GetDefaultPool())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> restored, + Split::Deserialize(bytes.data(), bytes.size(), GetDefaultPool())); + splits.push_back(std::move(restored)); + } + for (bool with_catalog : {false, true}) { + ReadContextBuilder builder(other->TablePath()); + if (with_catalog) { + builder.WithCatalog(other_catalog, Identifier("foo", "bar")); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(context))); + ASSERT_FALSE(splits.empty()); + ASSERT_NOK_WITH_MSG(read->CreateReader(splits.front()), "split was authorized for"); + ASSERT_NOK_WITH_MSG(read->CreateReader(splits), "split was authorized for"); + } + + // A worker can use an equivalent spelling of the authorized table's location. + for (const std::string& path : + {authorized->TablePath() + "/", "file://" + authorized->TablePath()}) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> context, + ReadContextBuilder(path).Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::StructArray> rows, + ReadAll(std::move(context), splits)); + AssertColumn(*rows, "phone", R"(["13800001111", "13900002222", null])"); + } +} + +TEST(TableTest, TestQueryAuthReadRejectsSplitsAuthorizedForAnotherBranch) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + const std::string path = helper->TablePath(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, "dev") + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, /*primary_keys=*/{}, + QueryAuthTableOptions("orc")) + .status()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, + PlanQueryAuth(path, catalog, /*predicate=*/nullptr, std::nullopt)); + ASSERT_FALSE(plan->Splits().empty()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> context, + ReadContextBuilder(path).WithBranch("dev").Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, TableRead::Create(std::move(context))); + ASSERT_NOK_WITH_MSG(read->CreateReader(plan->Splits().front()), + "split was authorized for table " + path + " (branch main)"); + ASSERT_NOK_WITH_MSG(read->CreateReader(plan->Splits()), + "but the read targets " + path + " (branch dev)"); + + auto planned = std::dynamic_pointer_cast<QueryAuthSplit>(plan->Splits().front()); + ASSERT_NE(nullptr, planned); + for (bool missing_path : {false, true}) { + QueryAuthScope incomplete_scope = planned->GetScope().value(); + if (missing_path) { + incomplete_scope.table_path.clear(); + } else { + incomplete_scope.branch.clear(); + } + ASSERT_NOK_WITH_MSG(read->CreateReader(std::make_shared<QueryAuthSplit>( + planned->GetSplit(), planned->GetAuthResult(), incomplete_scope)), + "does not record the table path and branch"); + } +} + +TEST(TableTest, TestCatalogDisablesQueryAuthOnlyForTheTableAtItsLocation) { + auto disabled_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(disabled_dir); + auto other_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(other_dir); + std::map<std::string, std::string> options = QueryAuthTableOptions("orc"); + options.erase(Options::QUERY_AUTH_ENABLED); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> disabled, + CreateQueryAuthTable(disabled_dir->Str(), /*primary_keys=*/{}, options, /*commits=*/{})); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> other, + CreateQueryAuthTable(other_dir->Str(), /*primary_keys=*/{}, options, /*commits=*/{})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> auth_catalog, + CreateQueryAuthCatalog(*disabled, disabled_dir->GetFileSystem())); + ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> disabled_schema, + disabled->LatestSchema()); + ASSERT_TRUE(disabled_schema); + auto plain_catalog = std::make_shared<MockVersionManagedCatalog>(); + plain_catalog->SetFileSystem(disabled_dir->GetFileSystem()); + plain_catalog->SetTableSchema(checked_pointer_cast<Schema>(disabled_schema.value())); + plain_catalog->SetTableLocation(disabled->TablePath()); + const Identifier identifier("foo", "bar"); + auto mismatch = [&](const std::string& operation) { + return fmt::format( + "Invalid: The {} targets {}, but the catalog keeps table foo.bar, whose query auth " + "settings it would apply, at {}", + operation, other->TablePath(), disabled->TablePath()); + }; + + for (const std::shared_ptr<Catalog>& catalog : + {std::shared_ptr<Catalog>(auth_catalog), std::shared_ptr<Catalog>(plain_catalog)}) { + for (const std::string& suffix : {std::string(), std::string("$files")}) { + SCOPED_TRACE(suffix); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(other->TablePath() + suffix) + .WithCatalog(catalog, identifier) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), mismatch("scan")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(other->TablePath() + suffix) + .WithCatalog(catalog, identifier) + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), mismatch("read")); + + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(disabled->TablePath() + suffix) + .WithCatalog(catalog, identifier) + .Finish()); + ASSERT_OK(TableScan::Create(std::move(scan_context)).status()); + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(disabled->TablePath() + suffix) + .WithCatalog(catalog, identifier) + .Finish()); + ASSERT_OK(TableRead::Create(std::move(read_context)).status()); + } + } + ASSERT_TRUE(auth_catalog->auth_identifiers.empty()); +} + +TEST(TableTest, TestQueryAuthTableRefusesWhatCannotApplyItsRules) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::optional<std::shared_ptr<TableSchema>> table_schema, + helper->LatestSchema()); + ASSERT_TRUE(table_schema); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + const Identifier identifier("foo", "bar"); + const std::string disabled = + "Invalid: Option 'query-auth.enabled' is not allowed to be disabled dynamically"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(path).Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "NotImplemented: Cannot read a table with 'query-auth.enabled=true' " + "without the catalog"); + auto plain_catalog = std::make_shared<MockVersionManagedCatalog>(); + plain_catalog->SetFileSystem(dir->GetFileSystem()); + plain_catalog->SetTableSchema(checked_pointer_cast<Schema>(table_schema.value())); + ASSERT_OK_AND_ASSIGN(scan_context, + ScanContextBuilder(path).WithCatalog(plain_catalog, identifier).Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "its catalog does not support the query auth interface"); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> disabling_scan_context, + ScanContextBuilder(path).AddOption(Options::QUERY_AUTH_ENABLED, "false").Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(disabling_scan_context)), disabled); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> disabling_read_context, + ReadContextBuilder(path).AddOption(Options::QUERY_AUTH_ENABLED, "false").Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(disabling_read_context)), disabled); + + auto explicit_read_schema = std::make_unique<::ArrowSchema>(); + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema({arrow::field("name", arrow::utf8())}), + explicit_read_schema.get()) + .ok()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> explicit_schema_read_context, + ReadContextBuilder(path).SetReadSchema(std::move(explicit_read_schema)).Finish()); + Status explicit_schema_status = + TableRead::Create(std::move(explicit_schema_read_context)).status(); + ASSERT_TRUE(explicit_schema_status.IsNotImplemented()) << explicit_schema_status.ToString(); + ASSERT_NOK_WITH_MSG(explicit_schema_status, + "An explicit read schema is not supported on a query-auth table"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> search_read_context, + ReadContextBuilder(path).EnablePreciseFileIndexEvaluate(true).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(search_read_context)), + "Search is not supported on a query-auth table"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<RealtimeContext> realtime_context, + RealtimeContext::Create()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> realtime_read_context, + ReadContextBuilder(path).WithRealtimeContext(realtime_context).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(realtime_read_context)), + "Real-time reads are not supported on a query-auth table"); + const std::string row_id = SpecialFields::RowId().Name(); + const std::string not_in_schema = + fmt::format("NotImplemented: Column '{}' is not a column of the table schema", row_id); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> row_id_read_context, + ReadContextBuilder(path).SetReadFieldNames({"id", row_id}).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(row_id_read_context)), not_in_schema); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> row_id_by_id_read_context, + ReadContextBuilder(path).SetReadFieldIds({0, SpecialFields::RowId().Id()}).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(row_id_by_id_read_context)), + fmt::format("NotImplemented: Field {} is not a column of the table schema", + SpecialFields::RowId().Id())); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> row_id_filter_read_context, + ReadContextBuilder(path) + .SetReadFieldNames({"id"}) + .SetPredicate(PredicateBuilder::IsNotNull(/*field_index=*/1, row_id, FieldType::BIGINT)) + .EnablePredicateFilter(true) + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(row_id_filter_read_context)), not_in_schema); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> row_id_unfiltered_read_context, + ReadContextBuilder(path) + .SetReadFieldNames({"id"}) + .SetPredicate(PredicateBuilder::IsNotNull(/*field_index=*/1, row_id, FieldType::BIGINT)) + .EnablePredicateFilter(false) + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(row_id_unfiltered_read_context)), + not_in_schema); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<MockQueryAuthCatalog> auth_catalog, + CreateQueryAuthCatalog(*helper, dir->GetFileSystem())); + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(auth_catalog, identifier) + .SetReadFieldNames({"id", row_id}) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), not_in_schema); + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(auth_catalog, identifier) + .SetPredicate(PredicateBuilder::IsNotNull( + /*field_index=*/4, row_id, FieldType::BIGINT)) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), not_in_schema); + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(auth_catalog, identifier) + .SetPartitionFilter({{{"gone", "1"}}}) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "NotImplemented: Column 'gone' is not a column of the table schema"); + ASSERT_OK_AND_ASSIGN( + scan_context, ScanContextBuilder(path) + .WithCatalog(auth_catalog, identifier) + .SetGlobalIndexResult(BitmapGlobalIndexResult::FromRanges({Range(0, 1)})) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "Search is not supported on a query-auth table"); + ASSERT_OK_AND_ASSIGN(scan_context, ScanContextBuilder(path) + .WithCatalog(auth_catalog, identifier) + .WithRealtimeContext(realtime_context) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), + "Real-time reads are not supported on a query-auth table"); + ASSERT_TRUE(auth_catalog->auth_identifiers.empty()); + const std::string global_index_refused = + "NotImplemented: Global index search is not supported on a table with " + "'query-auth.enabled=true'"; + ASSERT_NOK_WITH_MSG( + GlobalIndexScan::Create(path, /*snapshot_id=*/std::nullopt, /*partitions=*/std::nullopt, + /*options=*/{}, dir->GetFileSystem(), CreateDefaultExecutor(), + /*pool=*/nullptr), + global_index_refused); + ASSERT_NOK_WITH_MSG(GlobalIndexScan::Create(path, /*snapshot_id=*/std::nullopt, + /*partition_filters=*/std::shared_ptr<Predicate>(), + /*options=*/{}, dir->GetFileSystem(), + CreateDefaultExecutor(), /*pool=*/nullptr), + global_index_refused); + + for (const std::string system_table : {"files", "manifests", "binlog", "audit_log", "ro"}) { + ASSERT_NOK_WITH_MSG( + SystemTableLoader::LoadFromPath(dir->GetFileSystem(), path + "$" + system_table, {}), + fmt::format("NotImplemented: System table '{}' is not supported on a table with " + "'query-auth.enabled=true'", + system_table)); + } + ASSERT_OK( + SystemTableLoader::LoadFromPath(dir->GetFileSystem(), path + "$options", {}).status()); +} + +TEST(TableTest, TestQueryAuthRefusesFallbackReadsAcrossBranches) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, CreateQueryAuthTable(dir->Str())); + const std::string path = PathUtil::JoinPath(dir->Str(), "foo.db/bar"); + for (const std::string enabled : {"false", "true"}) { + const std::string branch = "fallback_" + enabled; + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, branch) + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, + /*primary_keys=*/{}, {{Options::QUERY_AUTH_ENABLED, enabled}}) + .status()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> context, + ReadContextBuilder(path).AddOption(Options::SCAN_FALLBACK_BRANCH, branch).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(context)), + "Fallback reads are not supported on a query-auth table"); + } + + const std::string fallback_path = PathUtil::JoinPath(dir->Str(), "fallback_tbl"); + auto fallback_schema = arrow::schema({arrow::field("id", arrow::int32(), /*nullable=*/false)}); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), fallback_path) + .CreateTable(fallback_schema, /*partition_keys=*/{}, /*primary_keys=*/{}, + /*options=*/{}) + .status()); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), fallback_path, "fallback") + .CreateTable(fallback_schema, /*partition_keys=*/{}, /*primary_keys=*/{}, + /*options=*/{{Options::QUERY_AUTH_ENABLED, "true"}}) + .status()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> fallback_read_context, + ReadContextBuilder(fallback_path) + .AddOption(Options::SCAN_FALLBACK_BRANCH, "fallback") + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(fallback_read_context)), + "Cannot read fallback branch 'fallback', which enables " + "'query-auth.enabled', of a table that does not"); +} + +TEST(TableTest, TestQueryAuthSystemTablesValidateTheBranchRead) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + const std::string path = dir->Str(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> main_schema, + SchemaManager(dir->GetFileSystem(), path) + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, + /*primary_keys=*/{}, /*options=*/{})); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, "dev") + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::QUERY_AUTH_ENABLED, "true"}}) + .status()); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), path, "restricted") + .CreateTable(QueryAuthTableSchema(), /*partition_keys=*/{}, /*primary_keys=*/{}, + /*options=*/{}) + .status()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> restricted_catalog_schema, + TableSchema::Create(/*schema_id=*/0, QueryAuthTableSchema(), + /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::QUERY_AUTH_ENABLED, "true"}})); + auto catalog = std::make_shared<MockQueryAuthCatalog>(); + catalog->table_location = path; + catalog->SetFileSystem(dir->GetFileSystem()); + catalog->SetTableSchema( + checked_pointer_cast<Schema>(std::shared_ptr<TableSchema>(std::move(main_schema)))); + catalog->object_schemas["bar$branch_restricted"] = checked_pointer_cast<Schema>( + std::shared_ptr<TableSchema>(std::move(restricted_catalog_schema))); + const Identifier identifier("foo", "bar"); + for (const std::string name : {"files", "manifests", "binlog", "audit_log", "ro"}) { + SCOPED_TRACE(name); + const std::string system_path = path + "$" + name; + const std::string refused = fmt::format( + "NotImplemented: System table '{}' is not supported on a table with " + "'query-auth.enabled=true'", + name); + ASSERT_NOK_WITH_MSG(SystemTableLoader::LoadFromPath(dir->GetFileSystem(), system_path, + {{Options::BRANCH, "dev"}}), + refused); + ASSERT_NOK_WITH_MSG( + SystemTableLoader::LoadFromPath(dir->GetFileSystem(), path + "$branch_dev$" + name, {}), + refused); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(system_path).AddOption(Options::BRANCH, "dev").Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), refused); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(system_path).WithBranch("dev").Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), refused); + + const std::string restricted_path = path + "$branch_restricted$" + name; + ASSERT_OK_AND_ASSIGN( + scan_context, + ScanContextBuilder(restricted_path).WithCatalog(catalog, identifier).Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(scan_context)), refused); + ASSERT_OK_AND_ASSIGN( + read_context, + ReadContextBuilder(restricted_path).WithCatalog(catalog, identifier).Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), refused); + ASSERT_EQ(Identifier("foo", "bar$branch_restricted"), + catalog->LoadTableSchemaIdentifiers().back()); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> main_read_context, + ReadContextBuilder(path + "$files").Finish()); + ASSERT_OK(TableRead::Create(std::move(main_read_context)).status()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> main_scan_context, + ScanContextBuilder(path + "$files").WithCatalog(catalog, identifier).Finish()); + ASSERT_OK(TableScan::Create(std::move(main_scan_context)).status()); + ASSERT_OK_AND_ASSIGN( + main_read_context, + ReadContextBuilder(path + "$files").WithCatalog(catalog, identifier).Finish()); + ASSERT_OK(TableRead::Create(std::move(main_read_context)).status()); + ASSERT_NOK_WITH_MSG( + SystemTableLoader::LoadFromPath(dir->GetFileSystem(), path + "$branch_dev$files", + {{Options::BRANCH, "other"}}), + "a system table read is aimed at one branch, but both 'other' and 'dev' were named"); +} + TEST(TableTest, TestNewWriteCreatesAFileStoreWrite) { auto dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); @@ -367,6 +1917,7 @@ catalog->SetFileSystem(dir->GetFileSystem()); // A catalog write and commit read the current schema from the catalog on the main branch. catalog->SetTableSchema(table->LatestSchema()); + catalog->SetTableLocation(dir->Str()); // `TableScan::Create()` and friends build the operator straight from the context, which // carries the catalog and identifier it resolves the table through.
diff --git a/src/paimon/rest/resource_paths.cpp b/src/paimon/rest/resource_paths.cpp index 75cfad2..a8e47f0 100644 --- a/src/paimon/rest/resource_paths.cpp +++ b/src/paimon/rest/resource_paths.cpp
@@ -88,4 +88,9 @@ return Table(database_name, table_name) + "/token"; } +std::string ResourcePaths::AuthTable(const std::string& database_name, + const std::string& table_name) const { + return Table(database_name, table_name) + "/auth"; +} + } // namespace paimon
diff --git a/src/paimon/rest/resource_paths.h b/src/paimon/rest/resource_paths.h index 0e7bf5f..80945e2 100644 --- a/src/paimon/rest/resource_paths.h +++ b/src/paimon/rest/resource_paths.h
@@ -51,6 +51,9 @@ /// Path of the temporary file system credentials of one table. std::string TableToken(const std::string& database_name, const std::string& table_name) const; + /// Endpoint returning the row filters and column masks a query of one table is subject to. + std::string AuthTable(const std::string& database_name, const std::string& table_name) const; + private: /// "/v1" or "/v1/{encoded prefix}". std::string base_;
diff --git a/src/paimon/rest/resource_paths_test.cpp b/src/paimon/rest/resource_paths_test.cpp index 5a6cd3e..56fecd9 100644 --- a/src/paimon/rest/resource_paths_test.cpp +++ b/src/paimon/rest/resource_paths_test.cpp
@@ -40,6 +40,7 @@ ASSERT_EQ("/v1/my+prefix/databases/db/tables/t1/snapshot", paths.TableSnapshot("db", "t1")); ASSERT_EQ("/v1/my+prefix/databases/db/tables/t1/token", paths.TableToken("db", "t1")); ASSERT_EQ("/v1/my+prefix/databases/db%231/tables/t+1/token", paths.TableToken("db#1", "t 1")); + ASSERT_EQ("/v1/my+prefix/databases/db/tables/t1/auth", paths.AuthTable("db", "t1")); } TEST(ResourcePathsTest, WithoutPrefix) { @@ -47,6 +48,7 @@ ASSERT_EQ("/v1/databases", paths.Databases()); ASSERT_EQ("/v1/tables/rename", paths.RenameTable()); ASSERT_EQ("/v1/databases/db/tables/t1/token", paths.TableToken("db", "t1")); + ASSERT_EQ("/v1/databases/db/tables/t1/auth", paths.AuthTable("db", "t1")); } } // namespace paimon::test
diff --git a/src/paimon/rest/rest_api.cpp b/src/paimon/rest/rest_api.cpp index 5862fb3..a47c168 100644 --- a/src/paimon/rest/rest_api.cpp +++ b/src/paimon/rest/rest_api.cpp
@@ -373,4 +373,18 @@ resource_paths_.TableToken(identifier.GetDatabaseName(), identifier.GetObjectName()), {}); } +Result<AuthTableQueryResponse> RestApi::AuthTableQuery( + const Identifier& identifier, const std::optional<std::vector<std::string>>& select) const { + AuthTableQueryRequest request(select); + PAIMON_ASSIGN_OR_RAISE(std::string body, request.ToJsonString()); + std::string path = + resource_paths_.AuthTable(identifier.GetDatabaseName(), identifier.GetObjectName()); + // Only reads the rules, so replaying it is safe. + PAIMON_ASSIGN_OR_RAISE(RestHttpClient::Response response, + Execute("POST", path, {}, body, /*retry_safe=*/true)); + AuthTableQueryResponse entity; + PAIMON_RETURN_NOT_OK(ParseResponseBody(response.body, path, &entity)); + return entity; +} + } // namespace paimon
diff --git a/src/paimon/rest/rest_api.h b/src/paimon/rest/rest_api.h index d420955..8c56a3a 100644 --- a/src/paimon/rest/rest_api.h +++ b/src/paimon/rest/rest_api.h
@@ -160,6 +160,13 @@ /// `RestErrorDetail`. Result<GetTableTokenResponse> LoadTableToken(const Identifier& identifier) const; + /// Loads the row filters and column masks a query of one table is subject to. + /// + /// @param identifier Table the query reads. + /// @param select Columns the query reads, or null for all of them. + Result<AuthTableQueryResponse> AuthTableQuery( + const Identifier& identifier, const std::optional<std::vector<std::string>>& select) const; + /// Maps a non-successful http response to a status: 404 becomes `NotExist`, 409 /// becomes `Exist`, 400 becomes `Invalid`, 501 becomes `NotImplemented` and the /// other codes become `IOError`. A redirect returned while `follow_redirects` is
diff --git a/src/paimon/rest/rest_catalog.cpp b/src/paimon/rest/rest_catalog.cpp index 273804d..17e1aa0 100644 --- a/src/paimon/rest/rest_catalog.cpp +++ b/src/paimon/rest/rest_catalog.cpp
@@ -537,6 +537,25 @@ return api_->LoadSnapshot(load_identifier); } +Result<std::shared_ptr<TableQueryAuthResult>> RestCatalog::AuthTableQuery( + const Identifier& identifier, const std::optional<std::vector<std::string>>& select) const { + PAIMON_RETURN_NOT_OK(CatalogUtils::CheckNotSystemTable(identifier, "authTable")); + PAIMON_ASSIGN_OR_RAISE(Identifier load_identifier, ToLoadIdentifier(identifier)); + // As in Java, an external table has no catalog to ask and is read without rules. + PAIMON_ASSIGN_OR_RAISE(GetTableResponse table, api_->GetTable(load_identifier)); + if (table.IsExternal()) { + return TableQueryAuthResult::Create(/*filter=*/{}, /*column_masking=*/{}); + } + PAIMON_ASSIGN_OR_RAISE(AuthTableQueryResponse response, + api_->AuthTableQuery(load_identifier, select)); + return TableQueryAuthResult::Create(response.GetFilter(), response.GetColumnMasking()); +} + +Result<std::shared_ptr<Schema>> RestCatalog::LoadTableSchemaAndLocation( + const Identifier& identifier, std::string* location) const { + return LoadTableSchema(identifier, /*table_id=*/nullptr, location); +} + Result<bool> RestCatalog::CommitSnapshot(const Identifier& identifier, const std::optional<std::string>& table_uuid, const std::optional<std::string>& base_snapshot_uuid,
diff --git a/src/paimon/rest/rest_catalog.h b/src/paimon/rest/rest_catalog.h index 2f57e91..f0942dc 100644 --- a/src/paimon/rest/rest_catalog.h +++ b/src/paimon/rest/rest_catalog.h
@@ -25,6 +25,7 @@ #include <vector> #include "paimon/catalog/catalog.h" +#include "paimon/core/catalog/query_auth_catalog.h" #include "paimon/core/catalog/version_managed_catalog.h" #include "paimon/logging.h" #include "paimon/rest/rest_api.h" @@ -41,7 +42,7 @@ /// A catalog backed by a REST catalog server. Metadata operations are delegated to /// `RestApi`; table data is accessed through the file system configured by the /// server-merged options. -class RestCatalog : public Catalog, public VersionManagedCatalog { +class RestCatalog : public Catalog, public VersionManagedCatalog, public QueryAuthCatalog { public: /// Creates the catalog: fetches and merges "/v1/config" from the server configured /// by `CatalogOptions::URI`, then builds the file system from the merged options. @@ -114,6 +115,18 @@ const Identifier& identifier, const std::optional<std::string>& table_uuid, const std::string& commit_user, const CoreOptions& options) override; + /// Loads the row filters and column masks of one query of a table from the server, mirroring + /// Java `RESTCatalog.authTableQuery`. A 404 is `NotExist`, a 403 an `IOError` carrying a + /// `RestErrorDetail` with the code, and a 501 `NotImplemented`. An external table, which Java + /// reads without asking for rules, has none: the server is not asked for them. + Result<std::shared_ptr<TableQueryAuthResult>> AuthTableQuery( + const Identifier& identifier, + const std::optional<std::vector<std::string>>& select) const override; + + /// Loads both from one `GetTable` response. + Result<std::shared_ptr<Schema>> LoadTableSchemaAndLocation( + const Identifier& identifier, std::string* location) const override; + /// Options merged with the server side config. const std::map<std::string, std::string>& GetOptions() const override;
diff --git a/src/paimon/rest/rest_catalog_test.cpp b/src/paimon/rest/rest_catalog_test.cpp index 147a908..ec600d0 100644 --- a/src/paimon/rest/rest_catalog_test.cpp +++ b/src/paimon/rest/rest_catalog_test.cpp
@@ -32,7 +32,9 @@ #include <vector> #include "arrow/api.h" +#include "arrow/array/concatenate.h" #include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" #include "fmt/format.h" #include "fmt/ranges.h" #include "gtest/gtest.h" @@ -41,24 +43,39 @@ #include "paimon/catalog_options.h" #include "paimon/commit_context.h" #include "paimon/commit_message.h" +#include "paimon/common/data/generic_row.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/catalog/commit_lock.h" #include "paimon/core/catalog/commit_table_request.h" +#include "paimon/core/catalog/query_auth_catalog.h" +#include "paimon/core/catalog/table_query_auth_result.h" #include "paimon/core/core_options.h" #include "paimon/core/partition/partition_statistics.h" +#include "paimon/core/schema/schema_manager.h" #include "paimon/core/schema/table_schema.h" +#include "paimon/core/table/system/global_system_tables.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/defs.h" #include "paimon/file_store_commit.h" #include "paimon/fs/file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/read_context.h" +#include "paimon/record_batch.h" #include "paimon/rest/mock_rest_server.h" #include "paimon/rest/rest_api.h" +#include "paimon/scan_context.h" #include "paimon/schema/schema.h" #include "paimon/table/format/format_table.h" +#include "paimon/table/source/plan.h" +#include "paimon/table/source/table_read.h" +#include "paimon/table/source/table_scan.h" +#include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/snapshot_test_helper.h" +#include "paimon/testing/utils/test_helper.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -80,8 +97,11 @@ std::string path; std::string id = "1"; std::map<std::string, std::string> token = {{"fs.oss.accessKeyId", "ak-1"}}; + bool is_external = false; }; std::map<std::string, std::map<std::string, TableData>> databases; + // Rebinds the table after its next GET response to exercise concurrent drop and recreate. + std::optional<TableData> recreate_table_after_get; // headers of the last request, with lower-cased names std::map<std::string, std::string> last_headers; // paths the data token endpoint was called at, in order @@ -114,6 +134,13 @@ std::optional<std::string> current_snapshot; bool null_snapshot_response = false; std::string last_snapshot_table; + std::vector<std::string> auth_filter; + std::map<std::string, std::string> auth_column_masking; + std::map<std::string, std::pair<std::vector<std::string>, std::map<std::string, std::string>>> + auth_rules_by_table; + std::vector<std::string> auth_requests; + std::vector<std::string> auth_request_bodies; + std::vector<int32_t> auth_error_codes; // guards all fields above: the handler runs on the server's accept thread while // tests seed and inspect the state std::mutex mutex; @@ -142,10 +169,10 @@ std::string TableResponseJson(const std::string& name, const MockCatalogState::TableData& table) { std::string id_member = table.id.empty() ? "" : fmt::format(R"("id":"{}",)", table.id); - return fmt::format( - R"({{{}"name":"{}","path":"{}","isExternal":false,"schemaId":{},"schema":{},)" - R"("owner":"owner1","updatedAt":123}})", - id_member, name, table.path, table.schema_id, table.schema_json); + return fmt::format(R"({{{}"name":"{}","path":"{}","isExternal":{},"schemaId":{},"schema":{},)" + R"("owner":"owner1","updatedAt":123}})", + id_member, name, table.path, table.is_external, table.schema_id, + table.schema_json); } // Serves `names` one item per page to exercise the pagination loop of the client. @@ -257,6 +284,45 @@ std::string remainder = rest.substr(databases_prefix.size()); size_t tables_pos = remainder.find("/tables"); + const std::string auth_suffix = "/auth"; + if (tables_pos != std::string::npos && request.method == "POST" && + remainder.size() > auth_suffix.size() && + remainder.compare(remainder.size() - auth_suffix.size(), auth_suffix.size(), auth_suffix) == + 0) { + state->auth_requests.push_back(request.path); + state->auth_request_bodies.push_back(request.body); + if (!state->auth_error_codes.empty()) { + const int32_t code = state->auth_error_codes.front(); + state->auth_error_codes.erase(state->auth_error_codes.begin()); + return MockError(code, "", "", "auth failed"); + } + Result<AuthTableQueryRequest> auth_request = + AuthTableQueryRequest::FromJsonString(request.body); + if (!auth_request.ok()) { + return MockError(400, "", "", auth_request.status().ToString()); + } + const std::optional<std::vector<std::string>>& select = auth_request.value().GetSelect(); + const std::string tables_infix = "/tables/"; + const size_t table_start = tables_pos + tables_infix.size(); + const std::string table = + remainder.substr(table_start, remainder.size() - auth_suffix.size() - table_start); + auto table_rules = state->auth_rules_by_table.find(table); + const std::vector<std::string>& filter = table_rules != state->auth_rules_by_table.end() + ? table_rules->second.first + : state->auth_filter; + const std::map<std::string, std::string>& masking = + table_rules != state->auth_rules_by_table.end() ? table_rules->second.second + : state->auth_column_masking; + std::map<std::string, std::string> column_masking; + for (const auto& [column, transform] : masking) { + if (!select || std::find(select->begin(), select->end(), column) != select->end()) { + column_masking.emplace(column, transform); + } + } + AuthTableQueryResponse response(filter, column_masking); + return JsonResponse(200, response.ToJsonString().value()); + } + const std::string token_suffix = "/token"; if (tables_pos != std::string::npos && remainder.size() > token_suffix.size() && remainder.compare(remainder.size() - token_suffix.size(), token_suffix.size(), @@ -500,7 +566,12 @@ if (state->get_table_delay_millis > 0) { std::this_thread::sleep_for(std::chrono::milliseconds(state->get_table_delay_millis)); } - return JsonResponse(200, TableResponseJson(table_name, table_iter->second)); + std::string body = TableResponseJson(table_name, table_iter->second); + if (state->recreate_table_after_get) { + table_iter->second = std::move(state->recreate_table_after_get.value()); + state->recreate_table_after_get.reset(); + } + return JsonResponse(200, body); } if (request.method == "DELETE") { tables.erase(table_iter); @@ -1347,6 +1418,12 @@ ASSERT_NOK_WITH_MSG(catalog->GetFormatTable(Identifier("db1", "only_partitions")), "every one of its columns"); + seed("query_auth", + R"({"type": "format-table", "file.format": "parquet", "query-auth.enabled": "true"})", + R"(["dt"])", default_fields); + ASSERT_NOK_WITH_MSG(catalog->GetFormatTable(Identifier("db1", "query_auth")), + "query-auth.enabled is not supported on a format table"); + // An option this library does not honour keeps its own status code, so a caller can still // tell "not supported yet" from "bad table". seed("catalog_partitions", @@ -1665,6 +1742,504 @@ } } +TEST_F(RestCatalogTest, ReadThroughTheCatalogAppliesItsQueryAuth) { + std::unique_ptr<UniqueTestDirectory> dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + const std::map<std::string, std::string> table_options = { + {Options::FILE_FORMAT, "orc"}, {Options::BUCKET, "1"}, {Options::BUCKET_KEY, "f0"}}; + auto schema = arrow::schema({arrow::field("f0", arrow::int32(), /*nullable=*/false), + arrow::field("f1", arrow::utf8())}); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), dir->Str()) + .CreateTable(schema, /*partition_keys=*/{}, /*primary_keys=*/{}, table_options) + .status()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> helper, + TestHelper::Create(dir->Str(), table_options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(schema->fields()), R"([[1, "a"], [2, "b"]])", + /*partition_map=*/{}, /*bucket=*/0, + /*row_kinds=*/{})); + ASSERT_OK( + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, std::nullopt).status()); + + MockCatalogState::TableData table_data; + table_data.schema_json = + R"({"fields": [{"id": 0, "name": "f0", "type": "INT NOT NULL"},)" + R"( {"id": 1, "name": "f1", "type": "STRING"}],)" + R"( "partitionKeys": [], "primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + table_data.path = dir->Str(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> rest_catalog, CreateRestCatalog()); + ASSERT_OK(rest_catalog->CreateDatabase("db1", {}, false)); + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"] = table_data; + state_->auth_filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0,)" + R"( "name": "f0", "type": "INT NOT NULL"}}, "function": "GREATER_THAN",)" + R"( "literals": [1]})"}; + state_->auth_column_masking = {{"f1", R"({"name": "CONCAT", "inputs": ["***"]})"}}; + } + std::shared_ptr<Catalog> catalog(std::move(rest_catalog)); + const Identifier identifier("db1", "t1"); + + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->get_table_requests = 0; + } + const QueryAuthCatalog* auth_catalog = AsQueryAuthCatalog(catalog); + ASSERT_NE(nullptr, auth_catalog); + std::string location; + ASSERT_OK(auth_catalog->LoadTableSchemaAndLocation(identifier, &location).status()); + ASSERT_EQ(dir->Str(), location); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ(1, state_->get_table_requests); + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(dir->Str()) + .WithCatalog(catalog, identifier) + .SetReadFieldNames({"f1"}) + .SetPredicate(PredicateBuilder::GreaterThan(/*field_index=*/0, "f0", FieldType::INT, + Literal(static_cast<int32_t>(0)))) + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ( + std::vector<std::string>({fmt::format("/v1/{}/databases/db1/tables/t1/auth", kPrefix)}), + state_->auth_requests); + ASSERT_EQ(1u, state_->auth_request_bodies.size()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest sent, + AuthTableQueryRequest::FromJsonString(state_->auth_request_bodies[0])); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{"f1", "f0"}), + sent.GetSelect()); + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(dir->Str()).AddOption(Options::QUERY_AUTH_ENABLED, "true").Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader, read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> chunks, + ReadResultCollector::CollectResult(std::move(reader))); + ASSERT_EQ(1, chunks->length()); + auto rows = checked_pointer_cast<arrow::StructArray>( + arrow::Concatenate(chunks->chunks(), arrow::default_memory_pool()).ValueOrDie()); + ASSERT_EQ(2, checked_pointer_cast<arrow::Int32Array>(rows->GetFieldByName("f0"))->Value(0)); + ASSERT_EQ("***", + checked_pointer_cast<arrow::StringArray>(rows->GetFieldByName("f1"))->GetString(0)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> disabling_read_context, + ReadContextBuilder(dir->Str()) + .WithCatalog(catalog, identifier) + .AddOption(Options::QUERY_AUTH_ENABLED, "false") + .Finish()); + ASSERT_NOK_WITH_MSG( + TableRead::Create(std::move(disabling_read_context)), + "Invalid: Option 'query-auth.enabled' is not allowed to be disabled dynamically"); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> files_scan_context, + ScanContextBuilder(dir->Str() + "$files") + .WithCatalog(catalog, Identifier("db1", "t1$files")) + .Finish()); + ASSERT_NOK_WITH_MSG(TableScan::Create(std::move(files_scan_context)), + "NotImplemented: System table 'files' is not supported on a table with " + "'query-auth.enabled=true'"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> options_scan_context, + ScanContextBuilder(dir->Str() + "$options") + .WithCatalog(catalog, Identifier("db1", "t1$options")) + .Finish()); + ASSERT_OK(TableScan::Create(std::move(options_scan_context)).status()); + + // A global system table identifier must not skip the catalog policy of a data table, + // including a table whose published schema does not enable query auth. + const Identifier system_identifier("sys", "tables"); + for (const std::string& suffix : {std::string(), std::string("$files")}) { + SCOPED_TRACE(suffix); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> mismatched_scan_context, + ScanContextBuilder(dir->Str() + suffix) + .WithCatalog(catalog, system_identifier) + .Finish()); + Status scan_status = TableScan::Create(std::move(mismatched_scan_context)).status(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> mismatched_read_context, + ReadContextBuilder(dir->Str() + suffix) + .WithCatalog(catalog, system_identifier) + .Finish()); + Status read_status = TableRead::Create(std::move(mismatched_read_context)).status(); + ASSERT_NOK_WITH_MSG(scan_status, "names a global system table, but the scan targets"); + ASSERT_NOK_WITH_MSG(read_status, "names a global system table, but the read targets"); + } + + const std::string global_path = PathUtil::JoinPath(dir->Str(), "sys/catalog_options"); + const Identifier global_identifier("sys", "catalog_options"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> global_scan_context, + ScanContextBuilder(global_path) + .WithCatalog(catalog, global_identifier) + .AddOption(CatalogOptionsSystemTable::kEnabledOption, "true") + .Finish()); + ASSERT_OK(TableScan::Create(std::move(global_scan_context)).status()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> global_read_context, + ReadContextBuilder(global_path) + .WithCatalog(catalog, global_identifier) + .AddOption(CatalogOptionsSystemTable::kEnabledOption, "true") + .Finish()); + ASSERT_OK(TableRead::Create(std::move(global_read_context)).status()); +} + +TEST_F(RestCatalogTest, GlobalPartitionsProtectsQueryAuthPartitionValues) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + std::map<std::string, std::string> table_options = { + {Options::FILE_FORMAT, "orc"}, {Options::BUCKET, "1"}, {Options::BUCKET_KEY, "v"}}; + auto schema = + arrow::schema({arrow::field("tenant", arrow::utf8()), arrow::field("v", arrow::int32())}); + SchemaManager schema_manager(dir->GetFileSystem(), dir->Str()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> initial_schema, + schema_manager.CreateTable(schema, /*partition_keys=*/{"tenant"}, + /*primary_keys=*/{}, table_options)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<TestHelper> helper, + TestHelper::Create(dir->Str(), table_options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(schema->fields()), R"([["secret", 1]])", + /*partition_map=*/{{"tenant", "secret"}}, /*bucket=*/0, + /*row_kinds=*/{})); + ASSERT_OK( + helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, std::nullopt).status()); + + MockCatalogState::TableData table_data; + ASSERT_OK_AND_ASSIGN(table_data.schema_json, initial_schema->ToJsonString()); + table_data.path = dir->Str(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> catalog, CreateRestCatalog()); + ASSERT_OK(catalog->CreateDatabase("db1", {}, /*ignore_if_exists=*/false)); + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"] = table_data; + } + GlobalSystemTableContext context; + context.catalog = catalog.get(); + context.fs = dir->GetFileSystem(); + PartitionsSystemTable partitions(context); + ASSERT_OK_AND_ASSIGN(std::vector<GenericRow> rows, partitions.BuildRows()); + ASSERT_EQ(1u, rows.size()); + ASSERT_EQ("tenant=secret", rows[0].GetString(2).ToString()); + + table_options[Options::QUERY_AUTH_ENABLED] = "true"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableSchema> enabled_schema, + TableSchema::Create(/*schema_id=*/1, schema, /*partition_keys=*/{"tenant"}, + /*primary_keys=*/{}, table_options)); + ASSERT_OK_AND_ASSIGN(std::string enabled_json, enabled_schema->ToJsonString()); + + // A disabled policy for the old location must not authorize the recreated table's location. + auto old_dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(old_dir); + MockCatalogState::TableData recreated = table_data; + recreated.schema_json = enabled_json; + recreated.id = "2"; + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"].path = old_dir->Str(); + state_->recreate_table_after_get = recreated; + } + ASSERT_OK_AND_ASSIGN(rows, partitions.BuildRows()); + ASSERT_TRUE(rows.empty()); + ASSERT_OK_AND_ASSIGN(rows, partitions.BuildRows()); + ASSERT_TRUE(rows.empty()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"] = table_data; + } + + // The published schema enables query auth before the catalog's schema catches up. + ASSERT_OK(dir->GetFileSystem()->WriteFile( + PathUtil::JoinPath(schema_manager.SchemaDirectory(), "schema-1"), enabled_json, + /*overwrite=*/false)); + ASSERT_OK_AND_ASSIGN(rows, partitions.BuildRows()); + ASSERT_TRUE(rows.empty()); +} + +TEST_F(RestCatalogTest, QueryAuthRejectsColumnsReplacedOnlyInTheCatalog) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto schema = + arrow::schema({arrow::field("f0", arrow::int32()), arrow::field("f1", arrow::utf8())}); + ASSERT_OK(SchemaManager(dir->GetFileSystem(), dir->Str()) + .CreateTable(schema, /*partition_keys=*/{}, /*primary_keys=*/{}, /*options=*/{}) + .status()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> rest_catalog, CreateRestCatalog()); + ASSERT_OK(rest_catalog->CreateDatabase("db1", {}, false)); + { + std::lock_guard<std::mutex> lock(state_->mutex); + MockCatalogState::TableData data; + data.path = dir->Str(); + data.schema_json = R"({"fields": [{"id": 0, "name": "f0", "type": "INT"},)" + R"({"id": 1, "name": "f1", "type": "STRING"}], "partitionKeys": [],)" + R"("primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + state_->databases["db1"]["t1"] = data; + } + std::shared_ptr<Catalog> catalog(std::move(rest_catalog)); + auto new_scan = + [&](const std::vector<std::string>& fields) -> Result<std::unique_ptr<TableScan>> { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ScanContext> context, + ScanContextBuilder(dir->Str()) + .WithCatalog(catalog, Identifier("db1", "t1")) + .SetReadFieldNames(fields) + .Finish()); + return TableScan::Create(std::move(context)); + }; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, new_scan({"f1"})); + ASSERT_OK(scan->CreatePlan().status()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"].schema_json = + R"({"fields": [{"id": 0, "name": "f0", "type": "INT"},)" + R"({"id": 9, "name": "f1", "type": "STRING"}], "partitionKeys": [],)" + R"("primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + } + const std::string replaced = + "Column 'f1' is field 1 in schema 0 the scan plans with, but " + "field 9 in catalog schema 0"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> stale_scan, new_scan({"f1"})); + ASSERT_NOK_WITH_MSG(stale_scan->CreatePlan(), replaced); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), replaced); + + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->auth_filter = { + R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 1,)" + R"("name": "f1", "type": "STRING"}}, "function": "IS_NOT_NULL", "literals": []})"}; + } + ASSERT_OK_AND_ASSIGN(scan, new_scan({"f0"})); + ASSERT_NOK_WITH_MSG(scan->CreatePlan(), "rules were bound to field 9"); +} + +TEST_F(RestCatalogTest, ScanOfABranchFollowsTheQueryAuthOfThatBranch) { + std::unique_ptr<UniqueTestDirectory> dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + const std::map<std::string, std::string> table_options = { + {Options::FILE_FORMAT, "orc"}, {Options::BUCKET, "1"}, {Options::BUCKET_KEY, "f0"}}; + auto schema = arrow::schema({arrow::field("f0", arrow::int32(), /*nullable=*/false), + arrow::field("f1", arrow::utf8())}); + for (const std::string& branch : {std::string("main"), std::string("dev")}) { + ASSERT_OK( + SchemaManager(dir->GetFileSystem(), dir->Str(), branch) + .CreateTable(schema, /*partition_keys=*/{}, /*primary_keys=*/{}, table_options) + .status()); + } + for (const auto& [branch, rows] : std::vector<std::pair<std::string, std::string>>{ + {"main", R"([[1, "main-a"], [2, "main-b"]])"}, + {"dev", R"([[11, "dev-a"], [12, "dev-b"], [13, "dev-c"]])"}}) { + std::map<std::string, std::string> options = table_options; + if (branch != "main") { + options[Options::BRANCH] = branch; + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestHelper> helper, + TestHelper::Create(dir->Str(), options, /*is_streaming_mode=*/false)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch, + TestHelper::MakeRecordBatch(arrow::struct_(schema->fields()), rows, + /*partition_map=*/{}, /*bucket=*/0, + /*row_kinds=*/{})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, std::nullopt) + .status()); + } + + MockCatalogState::TableData dev_data; + dev_data.schema_json = + R"({"fields": [{"id": 0, "name": "f0", "type": "INT NOT NULL"},)" + R"( {"id": 1, "name": "f1", "type": "STRING"}],)" + R"( "partitionKeys": [], "primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + dev_data.path = dir->Str(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> rest_catalog, CreateRestCatalog()); + ASSERT_OK(rest_catalog->CreateDatabase("db1", {}, false)); + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1$branch_dev"] = dev_data; + state_->auth_column_masking = {{"f1", R"({"name": "NULL"})"}}; + state_->auth_rules_by_table["t1$branch_dev"] = { + {R"({"kind": "LEAF", "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0,)" + R"( "name": "f0", "type": "INT NOT NULL"}}, "function": "GREATER_THAN",)" + R"( "literals": [11]})"}, + {{"f1", R"({"name": "SUBSTRING", "inputs": [{"index": 1, "name": "f1",)" + R"( "type": "STRING"}, 1, 3]})"}}}; + } + std::shared_ptr<Catalog> catalog(std::move(rest_catalog)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ScanContext> scan_context, + ScanContextBuilder(dir->Str()) + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableScan> scan, + TableScan::Create(std::move(scan_context))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<Plan> plan, scan->CreatePlan()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ(std::vector<std::string>( + {fmt::format("/v1/{}/databases/db1/tables/t1$branch_dev/auth", kPrefix)}), + state_->auth_requests); + } + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(dir->Str()) + .WithCatalog(catalog, Identifier("db1", "t1")) + .AddOption(Options::BRANCH, "dev") + .Finish()); + ASSERT_EQ(Identifier("db1", "t1$branch_dev"), read_context->GetIdentifier().value()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<TableRead> read, + TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<BatchReader> reader, read->CreateReader(plan->Splits())); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> chunks, + ReadResultCollector::CollectResult(std::move(reader))); + auto rows = checked_pointer_cast<arrow::StructArray>( + arrow::Concatenate(chunks->chunks(), arrow::default_memory_pool()).ValueOrDie()); + ASSERT_TRUE(rows->GetFieldByName("f0")->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[12, 13]").ValueOrDie())) + << rows->ToString(); + ASSERT_TRUE(rows->GetFieldByName("f1")->Equals( + *arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["dev", "dev"])") + .ValueOrDie())) + << rows->ToString(); +} + +TEST_F(RestCatalogTest, ReadRejectsAFallbackBranchProtectedOnlyByTheCatalog) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto schema = arrow::schema({arrow::field("f0", arrow::int32(), /*nullable=*/false)}); + for (const std::string branch : {"main", "dev"}) { + ASSERT_OK(SchemaManager(dir->GetFileSystem(), dir->Str(), branch) + .CreateTable(schema, /*partition_keys=*/{}, /*primary_keys=*/{}, + {{Options::FILE_FORMAT, "orc"}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "f0"}}) + .status()); + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> rest_catalog, CreateRestCatalog()); + ASSERT_OK(rest_catalog->CreateDatabase("db1", {}, false)); + { + std::lock_guard<std::mutex> lock(state_->mutex); + for (const auto& [table, enabled] : std::vector<std::pair<std::string, std::string>>{ + {"t1", "true"}, {"t1$branch_dev", "false"}}) { + MockCatalogState::TableData data; + data.path = dir->Str(); + data.schema_json = + R"({"fields": [{"id": 0, "name": "f0", "type": "INT NOT NULL"}],)" + R"("partitionKeys": [], "primaryKeys": [], "options": {"query-auth.enabled": ")" + + enabled + R"("}})"; + state_->databases["db1"][table] = data; + } + } + std::shared_ptr<Catalog> catalog(std::move(rest_catalog)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<ReadContext> read_context, + ReadContextBuilder(dir->Str()) + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch("dev") + .AddOption(Options::SCAN_FALLBACK_BRANCH, "main") + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), + "Cannot read fallback branch 'main', which enables 'query-auth.enabled'"); + for (const auto& [branch, message] : std::vector<std::pair<std::string, std::string>>{ + {"MAIN", "a catalog names branch 'MAIN' as it names the main branch"}, + {"dev$options", "a branch a catalog addresses cannot contain '$'"}}) { + ASSERT_OK_AND_ASSIGN(read_context, ReadContextBuilder(dir->Str()) + .WithCatalog(catalog, Identifier("db1", "t1")) + .WithBranch("dev") + .AddOption(Options::SCAN_FALLBACK_BRANCH, branch) + .Finish()); + ASSERT_NOK_WITH_MSG(TableRead::Create(std::move(read_context)), message); + } +} + +TEST_F(RestCatalogTest, AuthTableQueryParsesTheRules) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> catalog, CreateRestCatalog()); + MockCatalogState::TableData table_data; + table_data.schema_json = + R"({"fields": [{"id": 0, "name": "f1", "type": "STRING"}], "partitionKeys": [],)" + R"( "primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + table_data.path = "wh1/db1.db/t1"; + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1"] = table_data; + state_->databases["db1"]["t1$branch_b1"] = table_data; + state_->auth_column_masking = {{"f1", R"({"name": "NULL"})"}}; + } + ASSERT_OK_AND_ASSIGN( + std::shared_ptr<TableQueryAuthResult> result, + catalog->AuthTableQuery(Identifier("db1", "t1$branch_b1"), std::vector<std::string>{"f1"})); + ASSERT_EQ(std::set<std::string>({"f1"}), result->MaskedFields()); + ASSERT_EQ(nullptr, result->RowFilter()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ(fmt::format("/v1/{}/databases/db1/tables/t1$branch_b1/auth", kPrefix), + state_->auth_requests.back()); + state_->auth_column_masking = {{"f1", R"({"name": "UPPER", "inputs": []})"}}; + } + Status unknown = catalog->AuthTableQuery(Identifier("db1", "t1"), std::nullopt).status(); + ASSERT_TRUE(unknown.IsNotImplemented()) << unknown.ToString(); + ASSERT_NOK_WITH_MSG(catalog->AuthTableQuery(Identifier("db1", "t1$files"), std::nullopt), + "authTable"); + + table_data.is_external = true; + size_t auth_requests = 0; + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["external"] = table_data; + auth_requests = state_->auth_requests.size(); + } + ASSERT_OK_AND_ASSIGN(result, + catalog->AuthTableQuery(Identifier("db1", "external"), std::nullopt)); + ASSERT_FALSE(result->HasRules()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ(auth_requests, state_->auth_requests.size()); + } + ASSERT_NOK(catalog->AuthTableQuery(Identifier("db1", "missing"), std::nullopt)); +} + +TEST_F(RestCatalogTest, AuthTableQueryRetriesAnUnavailableServerButNotARefusal) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> catalog, + CreateRestCatalog(FastRetryConfig())); + MockCatalogState::TableData table_data; + table_data.schema_json = + R"({"fields": [{"id": 0, "name": "f1", "type": "STRING"}], "partitionKeys": [],)" + R"( "primaryKeys": [], "options": {"query-auth.enabled": "true"}})"; + table_data.path = "wh1/db1.db/t1"; + { + std::lock_guard<std::mutex> lock(state_->mutex); + state_->databases["db1"]["t1$branch_b1"] = table_data; + state_->auth_column_masking = {{"f1", R"({"name": "NULL"})"}}; + state_->auth_error_codes = {503}; + } + const Identifier identifier("db1", "t1$branch_b1"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<TableQueryAuthResult> result, + catalog->AuthTableQuery(identifier, std::vector<std::string>{"f1"})); + ASSERT_EQ(std::set<std::string>({"f1"}), result->MaskedFields()); + { + std::lock_guard<std::mutex> lock(state_->mutex); + const std::string path = + fmt::format("/v1/{}/databases/db1/tables/t1$branch_b1/auth", kPrefix); + ASSERT_EQ(std::vector<std::string>({path, path}), state_->auth_requests); + ASSERT_EQ(state_->auth_request_bodies[0], state_->auth_request_bodies[1]); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest sent, + AuthTableQueryRequest::FromJsonString(state_->auth_request_bodies[1])); + ASSERT_TRUE(sent.GetSelect()); + ASSERT_EQ(std::vector<std::string>({"f1"}), sent.GetSelect().value()); + state_->auth_error_codes = {403}; + } + ASSERT_NOK_WITH_MSG(catalog->AuthTableQuery(identifier, std::vector<std::string>{"f1"}), + "forbidden"); + { + std::lock_guard<std::mutex> lock(state_->mutex); + ASSERT_EQ(3u, state_->auth_requests.size()); + } +} + TEST_F(RestCatalogTest, CommitSnapshot) { ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestCatalog> catalog, CreateRestCatalog()); ASSERT_OK(catalog->CreateDatabase("db1", {}, false)); @@ -2434,4 +3009,70 @@ ASSERT_EQ(1, request_count.load()); } +TEST(RestApiTest, AuthTableQuery) { + std::mutex mutex; + std::vector<MockRestServer::Request> requests; + std::atomic<int32_t> response_code{200}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<MockRestServer> server, + MockRestServer::Start([&](const MockRestServer::Request& request) { + { + std::lock_guard<std::mutex> guard(mutex); + requests.push_back(request); + } + MockRestServer::Response response; + response.code = response_code.load(); + if (response.code == 200) { + response.body = + R"({"filter": ["f1", "f2"], "columnMasking": {"name": "m1"}})"; + } else { + response.body = fmt::format( + R"({{"message": "refused", "code": {}}})", response.code); + } + return response; + })); + std::map<std::string, std::string> options = { + {CatalogOptions::URI, server->GetBaseUri()}, + {CatalogOptions::TOKEN_PROVIDER, "bear"}, + {CatalogOptions::TOKEN, kToken}, + {RestApi::kOptionUrlPrefix, "p1"}, + }; + ASSERT_OK_AND_ASSIGN(std::unique_ptr<RestApi> api, + RestApi::Create(options, "", /*config_required=*/false)); + + ASSERT_OK_AND_ASSIGN( + AuthTableQueryResponse response, + api->AuthTableQuery(Identifier("db1", "t1$branch_b1"), std::vector<std::string>{"id"})); + ASSERT_EQ(std::vector<std::string>({"f1", "f2"}), response.GetFilter()); + ASSERT_EQ((std::map<std::string, std::string>{{"name", "m1"}}), response.GetColumnMasking()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryResponse all_columns, + api->AuthTableQuery(Identifier("db1", "t1"), std::nullopt)); + ASSERT_EQ(response.GetFilter(), all_columns.GetFilter()); + { + std::lock_guard<std::mutex> guard(mutex); + ASSERT_EQ(2u, requests.size()); + ASSERT_EQ("POST", requests[0].method); + ASSERT_EQ("/v1/p1/databases/db1/tables/t1$branch_b1/auth", requests[0].path); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest sent, + AuthTableQueryRequest::FromJsonString(requests[0].body)); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{"id"}), + sent.GetSelect()); + ASSERT_EQ("/v1/p1/databases/db1/tables/t1/auth", requests[1].path); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest sent_all_columns, + AuthTableQueryRequest::FromJsonString(requests[1].body)); + ASSERT_FALSE(sent_all_columns.GetSelect()); + } + + response_code = 404; + Status not_exist = api->AuthTableQuery(Identifier("db1", "t1"), std::nullopt).status(); + ASSERT_TRUE(not_exist.IsNotExist()) << not_exist.ToString(); + response_code = 403; + Status forbidden = api->AuthTableQuery(Identifier("db1", "t1"), std::nullopt).status(); + ASSERT_TRUE(forbidden.IsIOError()) << forbidden.ToString(); + ASSERT_NE(nullptr, forbidden.detail()); + ASSERT_EQ(403, static_cast<const RestErrorDetail*>(forbidden.detail().get())->GetCode()); + response_code = 501; + Status not_implemented = api->AuthTableQuery(Identifier("db1", "t1"), std::nullopt).status(); + ASSERT_TRUE(not_implemented.IsNotImplemented()) << not_implemented.ToString(); +} + } // namespace paimon::test
diff --git a/src/paimon/rest/rest_messages.cpp b/src/paimon/rest/rest_messages.cpp index 177b26d..aa661f2 100644 --- a/src/paimon/rest/rest_messages.cpp +++ b/src/paimon/rest/rest_messages.cpp
@@ -62,6 +62,9 @@ constexpr const char kFieldCommitUser[] = "commitUser"; constexpr const char kFieldAcquired[] = "acquired"; constexpr const char kFieldLeaseMillis[] = "leaseMillis"; +constexpr const char kFieldSelect[] = "select"; +constexpr const char kFieldFilter[] = "filter"; +constexpr const char kFieldColumnMasking[] = "columnMasking"; void AddOptionalStringMember(rapidjson::Value* obj, const char* key, const std::optional<std::string>& value, @@ -498,4 +501,47 @@ snapshot_ = RapidJsonUtil::DeserializeKeyValue<std::optional<Snapshot>>(obj, kFieldSnapshot); } +rapidjson::Value AuthTableQueryRequest::ToJson(rapidjson::Document::AllocatorType* allocator) const + noexcept(false) { + rapidjson::Value obj(rapidjson::kObjectType); + if (select_) { + obj.AddMember(rapidjson::StringRef(kFieldSelect), + RapidJsonUtil::SerializeValue(select_.value(), allocator).Move(), *allocator); + } + return obj; +} + +void AuthTableQueryRequest::FromJson(const rapidjson::Value& obj) noexcept(false) { + select_ = RapidJsonUtil::DeserializeKeyValue<std::optional<std::vector<std::string>>>( + obj, kFieldSelect); +} + +rapidjson::Value AuthTableQueryResponse::ToJson(rapidjson::Document::AllocatorType* allocator) const + noexcept(false) { + rapidjson::Value obj(rapidjson::kObjectType); + if (!filter_.empty()) { + obj.AddMember(rapidjson::StringRef(kFieldFilter), + RapidJsonUtil::SerializeValue(filter_, allocator).Move(), *allocator); + } + if (!column_masking_.empty()) { + obj.AddMember(rapidjson::StringRef(kFieldColumnMasking), + RapidJsonUtil::SerializeValue(column_masking_, allocator).Move(), *allocator); + } + return obj; +} + +void AuthTableQueryResponse::FromJson(const rapidjson::Value& obj) noexcept(false) { + // A member named twice fails rather than being read as either of its values. + Status unique = RapidJsonUtil::CheckNoDuplicateMembers(obj); + if (!unique.ok()) { + throw std::invalid_argument(unique.message()); + } + // A member of the wrong type fails rather than reading as no rules, which would let a query + // through unrestricted. + filter_ = RapidJsonUtil::DeserializeKeyValue<std::vector<std::string>>( + obj, kFieldFilter, std::vector<std::string>()); + column_masking_ = RapidJsonUtil::DeserializeKeyValue<std::map<std::string, std::string>>( + obj, kFieldColumnMasking, std::map<std::string, std::string>()); +} + } // namespace paimon
diff --git a/src/paimon/rest/rest_messages.h b/src/paimon/rest/rest_messages.h index 1f01cb8..fa2df96 100644 --- a/src/paimon/rest/rest_messages.h +++ b/src/paimon/rest/rest_messages.h
@@ -532,4 +532,53 @@ std::optional<Snapshot> snapshot_; }; +/// Asks for the row filters and column masks a query of one table is subject to. +class AuthTableQueryRequest : public Jsonizable<AuthTableQueryRequest> { + public: + /// @param select Columns the query reads, or null for all of them. + explicit AuthTableQueryRequest(const std::optional<std::vector<std::string>>& select) + : select_(select) {} + + rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const + noexcept(false) override; + void FromJson(const rapidjson::Value& obj) noexcept(false) override; + + const std::optional<std::vector<std::string>>& GetSelect() const { + return select_; + } + + AuthTableQueryRequest() = default; + + private: + std::optional<std::vector<std::string>> select_; +}; + +/// The row filters and column masks a query of one table is subject to. Both are kept as the +/// JSON the server sent: `filter` holds serialized predicates that all apply, and +/// `columnMasking` maps a column to the serialized transform that replaces its values. An absent +/// member is read as no rules of that kind. +class AuthTableQueryResponse : public Jsonizable<AuthTableQueryResponse> { + public: + AuthTableQueryResponse(const std::vector<std::string>& filter, + const std::map<std::string, std::string>& column_masking) + : filter_(filter), column_masking_(column_masking) {} + + rapidjson::Value ToJson(rapidjson::Document::AllocatorType* allocator) const + noexcept(false) override; + void FromJson(const rapidjson::Value& obj) noexcept(false) override; + + const std::vector<std::string>& GetFilter() const { + return filter_; + } + const std::map<std::string, std::string>& GetColumnMasking() const { + return column_masking_; + } + + AuthTableQueryResponse() = default; + + private: + std::vector<std::string> filter_; + std::map<std::string, std::string> column_masking_; +}; + } // namespace paimon
diff --git a/src/paimon/rest/rest_messages_test.cpp b/src/paimon/rest/rest_messages_test.cpp index 9aba13b..d030f53 100644 --- a/src/paimon/rest/rest_messages_test.cpp +++ b/src/paimon/rest/rest_messages_test.cpp
@@ -481,6 +481,95 @@ ASSERT_NOK_WITH_MSG(CommitLockResponse::FromJsonString("[]"), "value must be an object"); } +TEST(RestMessagesTest, AuthTableQueryRequestRoundTrip) { + AuthTableQueryRequest request(std::vector<std::string>{"id", "name"}); + ASSERT_OK_AND_ASSIGN(std::string json, request.ToJsonString()); + rapidjson::Document doc; + doc.Parse(json.c_str()); + ASSERT_FALSE(doc.HasParseError()); + ASSERT_EQ(1u, doc.MemberCount()); + ASSERT_TRUE(doc["select"].IsArray()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest parsed, AuthTableQueryRequest::FromJsonString(json)); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{"id", "name"}), + parsed.GetSelect()); + + AuthTableQueryRequest all_columns(std::nullopt); + ASSERT_OK_AND_ASSIGN(std::string all_columns_json, all_columns.ToJsonString()); + ASSERT_EQ("{}", all_columns_json); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest parsed_all_columns, + AuthTableQueryRequest::FromJsonString(all_columns_json)); + ASSERT_FALSE(parsed_all_columns.GetSelect()); + AuthTableQueryRequest no_columns(std::vector<std::string>{}); + ASSERT_OK_AND_ASSIGN(std::string no_columns_json, no_columns.ToJsonString()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest parsed_no_columns, + AuthTableQueryRequest::FromJsonString(no_columns_json)); + ASSERT_EQ(std::optional<std::vector<std::string>>(std::vector<std::string>{}), + parsed_no_columns.GetSelect()); + + const std::string column("name\0secret", 11); + AuthTableQueryRequest nul_column(std::vector<std::string>{"name", column}); + ASSERT_OK_AND_ASSIGN(std::string nul_column_json, nul_column.ToJsonString()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryRequest parsed_nul_column, + AuthTableQueryRequest::FromJsonString(nul_column_json)); + ASSERT_EQ(nul_column.GetSelect(), parsed_nul_column.GetSelect()); +} + +TEST(RestMessagesTest, AuthTableQueryResponseRoundTrip) { + std::string json = + R"({"filter": ["{\"kind\":\"LEAF\"}"], "columnMasking": {"name": "{\"name\":\"NULL\"}"},)" + R"( "unknown": 1})"; + ASSERT_OK_AND_ASSIGN(AuthTableQueryResponse response, + AuthTableQueryResponse::FromJsonString(json)); + ASSERT_EQ(std::vector<std::string>({R"({"kind":"LEAF"})"}), response.GetFilter()); + ASSERT_EQ((std::map<std::string, std::string>{{"name", R"({"name":"NULL"})"}}), + response.GetColumnMasking()); + + ASSERT_OK_AND_ASSIGN(std::string written, response.ToJsonString()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryResponse round_tripped, + AuthTableQueryResponse::FromJsonString(written)); + ASSERT_EQ(response.GetFilter(), round_tripped.GetFilter()); + ASSERT_EQ(response.GetColumnMasking(), round_tripped.GetColumnMasking()); + + const std::string column("name\0secret", 11); + AuthTableQueryResponse nul_column({}, {{"name", R"({"name": "NULL"})"}, + {column, R"({"name": "CONCAT", "inputs": ["***"]})"}}); + ASSERT_OK_AND_ASSIGN(std::string nul_column_json, nul_column.ToJsonString()); + ASSERT_OK_AND_ASSIGN(AuthTableQueryResponse parsed_nul_column, + AuthTableQueryResponse::FromJsonString(nul_column_json)); + ASSERT_EQ(nul_column.GetColumnMasking(), parsed_nul_column.GetColumnMasking()); + + for (const char* no_rules : {"{}", R"({"filter": null, "columnMasking": null})"}) { + ASSERT_OK_AND_ASSIGN(AuthTableQueryResponse empty, + AuthTableQueryResponse::FromJsonString(no_rules)); + ASSERT_TRUE(empty.GetFilter().empty()); + ASSERT_TRUE(empty.GetColumnMasking().empty()); + } + + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString(R"({"filter": "f"})"), + "value must be an array"); + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString(R"({"filter": [1]})"), + "value must be string"); + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString(R"({"columnMasking": ["m"]})"), + "value must be an object"); + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString(R"({"columnMasking": {"c": 1}})"), + "value must be string"); + + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString( + R"({"filter": [], "filter": ["{\"kind\":\"LEAF\"}"]})"), + "JSON object names member 'filter' more than once"); + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString( + R"({"columnMasking": {"c": "{\"name\":\"FIELD_REF\"}",)" + R"( "c": "{\"name\":\"NULL\"}"}})"), + "JSON object names member 'c' more than once"); + ASSERT_NOK_WITH_MSG( + AuthTableQueryResponse::FromJsonString(R"({"columnMasking": {}, "columnMasking": null})"), + "JSON object names member 'columnMasking' more than once"); + for (const std::string& suffix : {std::string(1, '\0'), std::string(1, '\0') + "not JSON"}) { + ASSERT_NOK_WITH_MSG(AuthTableQueryResponse::FromJsonString("{}" + suffix), + "JSON contains an unescaped NUL byte"); + } +} + TEST(RestMessagesTest, CreateDatabaseRequestRoundTrip) { CreateDatabaseRequest request("db1", {{"k1", "v1"}}); ASSERT_OK_AND_ASSIGN(std::string json, request.ToJsonString());
diff --git a/src/paimon/testing/mock/mock_catalog.h b/src/paimon/testing/mock/mock_catalog.h index 1a2f0b7..a24e246 100644 --- a/src/paimon/testing/mock/mock_catalog.h +++ b/src/paimon/testing/mock/mock_catalog.h
@@ -113,8 +113,8 @@ Result<std::shared_ptr<Table>> GetTable(const Identifier& identifier) const override { ++get_table_calls_; // Serves the per-table file system when one was set, matching `GetTableFileSystem` - // without recording the request, so a test can still tell the two apart. No location: - // this mock's `GetTableLocation` is the unsupported stub. + // without recording the request, so a test can still tell the two apart. No location, + // even when `SetTableLocation()` set one. return std::make_shared<Table>( table_file_system_ != nullptr ? table_file_system_ : file_system_, /*root_path=*/"", table_schema_, identifier, table_uuid_); @@ -235,8 +235,16 @@ Result<std::string> GetDatabaseLocation(const std::string&) const override { return Unsupported(); } + /// Serves the location set by `SetTableLocation()` for every identifier, or fails when none + /// was set. Result<std::string> GetTableLocation(const Identifier&) const override { - return Unsupported(); + if (!table_location_) { + return Unsupported(); + } + return table_location_.value(); + } + void SetTableLocation(const std::string& location) { + table_location_ = location; } std::string GetRootPath() const override { return ""; @@ -294,6 +302,7 @@ Result<bool> answer_; std::string table_uuid_; std::shared_ptr<Schema> table_schema_; + std::optional<std::string> table_location_; std::shared_ptr<FileSystem> file_system_; std::shared_ptr<FileSystem> table_file_system_; mutable std::vector<Identifier> table_file_system_requests_;
diff --git a/src/paimon/testing/utils/test_helper.h b/src/paimon/testing/utils/test_helper.h index 2bc1e9b..bc87fcf 100644 --- a/src/paimon/testing/utils/test_helper.h +++ b/src/paimon/testing/utils/test_helper.h
@@ -416,6 +416,10 @@ return channel_file_count; } + const std::string& TablePath() const { + return table_path_; + } + private: void CheckExternalPath(const std::vector<std::shared_ptr<CommitMessage>>& actuals) { for (const auto& actual : actuals) {
diff --git a/test/inte/blob_table_inte_test.cpp b/test/inte/blob_table_inte_test.cpp index fd706e6..1bd320e 100644 --- a/test/inte/blob_table_inte_test.cpp +++ b/test/inte/blob_table_inte_test.cpp
@@ -4719,8 +4719,10 @@ // reading fails because the upstream table does not exist at the catalog-issued location ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); - ASSERT_NOK_WITH_MSG(ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr), - "not found latest schema"); + Status missing_upstream = + ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr).status(); + ASSERT_TRUE(missing_upstream.IsNotExist()) << missing_upstream.ToString(); + ASSERT_NOK_WITH_MSG(missing_upstream, "nonexistent_table"); } TEST_P(BlobTableInteTest, TestBlobViewSkipsDanglingReferenceOfDeletedRow) { @@ -4738,8 +4740,10 @@ // without the deletion vector the dangling reference fails the read ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); - ASSERT_NOK_WITH_MSG(ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr), - "not found latest schema"); + Status missing_upstream = + ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr).status(); + ASSERT_TRUE(missing_upstream.IsNotExist()) << missing_upstream.ToString(); + ASSERT_NOK_WITH_MSG(missing_upstream, "nonexistent_table"); ASSERT_OK_AND_ASSIGN(std::string anchor_file_name, DeletionVectorTestHelper::RetrieveAnchorFileName(commit_msgs)); @@ -4782,8 +4786,10 @@ // without deletion vectors both groups' dangling references fail the read ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); - ASSERT_NOK_WITH_MSG(ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr), - "not found latest schema"); + Status missing_upstream = + ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr).status(); + ASSERT_TRUE(missing_upstream.IsNotExist()) << missing_upstream.ToString(); + ASSERT_NOK_WITH_MSG(missing_upstream, "nonexistent_table"); // each group's anchor gets a deletion vector dropping its row 0, the dangling one ASSERT_OK_AND_ASSIGN(std::string anchor_file_name0, @@ -4819,16 +4825,20 @@ // reading the whole split fails on the dangling reference of row 0 ASSERT_OK_AND_ASSIGN(auto plan, ScanTable(table_path)); - ASSERT_NOK_WITH_MSG(ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr), - "not found latest schema"); + Status missing_upstream = + ReadTable(table_path, schema->field_names(), plan, /*predicate=*/nullptr).status(); + ASSERT_TRUE(missing_upstream.IsNotExist()) << missing_upstream.ToString(); + ASSERT_NOK_WITH_MSG(missing_upstream, "nonexistent_table"); // a predicate keeping only row 1 does not rescue the query: it is evaluated above the split // read, while the pre-read takes none, so row 0's reference is still collected and resolved auto keep_row_1 = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT, Literal(101)); ASSERT_OK_AND_ASSIGN(auto filtered_plan, ScanTable(table_path, keep_row_1)); - ASSERT_NOK_WITH_MSG(ReadTable(table_path, schema->field_names(), filtered_plan, keep_row_1), - "not found latest schema"); + missing_upstream = + ReadTable(table_path, schema->field_names(), filtered_plan, keep_row_1).status(); + ASSERT_TRUE(missing_upstream.IsNotExist()) << missing_upstream.ToString(); + ASSERT_NOK_WITH_MSG(missing_upstream, "nonexistent_table"); // dropping the same row through the row ranges does: the pre-read honors them the same way // the main read does, so row 0 is never resolved
diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index ca88934..7bdb500 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp
@@ -4548,6 +4548,26 @@ /*expected_commit_messages=*/std::nullopt)); ASSERT_OK_AND_ASSIGN(auto catalog, Catalog::Create(warehouse, options)); + std::map<std::string, std::string> query_auth_options = options; + query_auth_options[Options::QUERY_AUTH_ENABLED] = "true"; + const Identifier query_auth_table("foo", "query_auth"); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); + ASSERT_OK(catalog->CreateTable(query_auth_table, &c_schema, /*partition_keys=*/{"dt", "region"}, + /*primary_keys=*/{}, query_auth_options, + /*ignore_if_exists=*/false)); + ArrowSchemaRelease(&c_schema); + ASSERT_OK_AND_ASSIGN(std::string query_auth_path, catalog->GetTableLocation(query_auth_table)); + ASSERT_OK_AND_ASSIGN(auto query_auth_helper, + TestHelper::Create(query_auth_path, query_auth_options, + /*is_streaming_mode=*/true)); + ASSERT_OK_AND_ASSIGN(batch, TestHelper::MakeRecordBatch( + arrow::struct_(fields), R"([["2026-07-13", "secret", 2]])", + /*partition_map=*/{{"dt", "2026-07-13"}, {"region", "secret"}}, + /*bucket=*/0, {})); + ASSERT_OK(query_auth_helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0, + /*expected_commit_messages=*/std::nullopt)); + auto fs = catalog->GetFileSystem(); ASSERT_OK_AND_ASSIGN(auto part_result, ReadGlobalSystemTable("partitions", catalog.get(), fs, warehouse, options));
diff --git a/test/test_data/query_auth/QueryAuthSplitFixtures.java b/test/test_data/query_auth/QueryAuthSplitFixtures.java new file mode 100644 index 0000000..c431e1f --- /dev/null +++ b/test/test_data/query_auth/QueryAuthSplitFixtures.java
@@ -0,0 +1,163 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import org.apache.paimon.catalog.TableQueryAuthResult; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.Timestamp; +import org.apache.paimon.globalindex.IndexedSplit; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.manifest.FileSource; +import org.apache.paimon.predicate.ConcatTransform; +import org.apache.paimon.predicate.FieldRef; +import org.apache.paimon.predicate.NullTransform; +import org.apache.paimon.predicate.PredicateBuilder; +import org.apache.paimon.predicate.SubstringTransform; +import org.apache.paimon.stats.SimpleStats; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.QueryAuthSplit; +import org.apache.paimon.table.source.Split; +import org.apache.paimon.table.source.SplitSerializer; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; +import org.apache.paimon.utils.JsonSerdeUtil; +import org.apache.paimon.utils.Range; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Writes the Java {@code SplitSerializer} bytes of the {@code QueryAuthSplit} fixtures read by + * {@code query_auth_split_test.cpp}. The checked-in fixtures were written by paimon-core + * 2.2-SNAPSHOT built from apache/paimon commit c2e001b7f2c0214cc501d23a89be44ae031b447d: + * + * <pre> + * javac -cp "$PAIMON_CLASSPATH" -d classes QueryAuthSplitFixtures.java + * java -cp "classes:$PAIMON_CLASSPATH" QueryAuthSplitFixtures test/test_data/query_auth + * </pre> + */ +public class QueryAuthSplitFixtures { + + public static void main(String[] args) throws Exception { + Path dir = Paths.get(args[0]); + RowType rowType = + RowType.of( + new org.apache.paimon.types.DataType[] { + DataTypes.INT().notNull(), + DataTypes.STRING(), + DataTypes.STRING(), + DataTypes.BIGINT() + }, + new String[] {"id", "name", "phone", "salary"}); + PredicateBuilder builder = new PredicateBuilder(rowType); + List<String> filter = + Arrays.asList( + JsonSerdeUtil.toFlatJson(builder.greaterThan(0, 1)), + JsonSerdeUtil.toFlatJson( + builder.notEqual(1, BinaryString.fromString("测试")))); + Map<String, String> masking = new HashMap<>(); + masking.put( + "name", + JsonSerdeUtil.toFlatJson( + new ConcatTransform( + Collections.singletonList(BinaryString.fromString("名字"))))); + masking.put( + "phone", + JsonSerdeUtil.toFlatJson( + new SubstringTransform( + Arrays.asList( + new FieldRef(2, "phone", DataTypes.STRING()), 1, 3)))); + masking.put("salary", JsonSerdeUtil.toFlatJson(NullTransform.INSTANCE)); + TableQueryAuthResult rules = new TableQueryAuthResult(filter, masking); + + DataFileMeta file = + DataFileMeta.create( + "file1.orc", + 100L, + 200L, + BinaryRow.EMPTY_ROW, + BinaryRow.EMPTY_ROW, + SimpleStats.EMPTY_STATS, + SimpleStats.EMPTY_STATS, + 50L, + 249L, + 0L, + 0, + Collections.emptyList(), + Timestamp.fromEpochMillis(1765535214349L), + 0L, + null, + FileSource.APPEND, + null, + null, + 50L, + null, + null); + DataSplit dataSplit = + DataSplit.builder() + .withSnapshot(1) + .withPartition(BinaryRow.EMPTY_ROW) + .withBucket(0) + .withBucketPath("data/test_table/bucket-0") + .withDataFiles(Collections.singletonList(file)) + .isStreaming(false) + .rawConvertible(true) + .build(); + IndexedSplit indexedSplit = + new IndexedSplit( + dataSplit, + Collections.singletonList(new Range(55, 56)), + new float[] {0.5f}); + + write(dir, "query_auth_split-data_split", new QueryAuthSplit(dataSplit, rules)); + write(dir, "query_auth_split-indexed_split", new QueryAuthSplit(indexedSplit, rules)); + write( + dir, + "query_auth_split-empty_rules", + new QueryAuthSplit( + dataSplit, + new TableQueryAuthResult( + Collections.emptyList(), Collections.emptyMap()))); + write( + dir, + "query_auth_split-null_rules", + new QueryAuthSplit(dataSplit, new TableQueryAuthResult(null, null))); + write(dir, "query_auth_split-null_auth_result", new QueryAuthSplit(dataSplit, null)); + + for (String json : filter) { + System.out.println("filter: " + json); + } + for (Map.Entry<String, String> mask : masking.entrySet()) { + System.out.println("mask " + mask.getKey() + ": " + mask.getValue()); + } + } + + private static void write(Path dir, String name, Split split) throws Exception { + byte[] bytes = SplitSerializer.serialize(split); + if (!Arrays.equals(bytes, SplitSerializer.serialize(SplitSerializer.deserialize(bytes)))) { + throw new IllegalStateException(name + " does not round-trip in Java"); + } + Files.write(dir.resolve(name), bytes); + } +}
diff --git a/test/test_data/query_auth/query_auth_split-data_split b/test/test_data/query_auth/query_auth_split-data_split new file mode 100644 index 0000000..73ee6c3 --- /dev/null +++ b/test/test_data/query_auth/query_auth_split-data_split Binary files differ
diff --git a/test/test_data/query_auth/query_auth_split-empty_rules b/test/test_data/query_auth/query_auth_split-empty_rules new file mode 100644 index 0000000..c48b40b --- /dev/null +++ b/test/test_data/query_auth/query_auth_split-empty_rules Binary files differ
diff --git a/test/test_data/query_auth/query_auth_split-indexed_split b/test/test_data/query_auth/query_auth_split-indexed_split new file mode 100644 index 0000000..7e50797 --- /dev/null +++ b/test/test_data/query_auth/query_auth_split-indexed_split Binary files differ
diff --git a/test/test_data/query_auth/query_auth_split-null_auth_result b/test/test_data/query_auth/query_auth_split-null_auth_result new file mode 100644 index 0000000..69b18de --- /dev/null +++ b/test/test_data/query_auth/query_auth_split-null_auth_result Binary files differ
diff --git a/test/test_data/query_auth/query_auth_split-null_rules b/test/test_data/query_auth/query_auth_split-null_rules new file mode 100644 index 0000000..bd5e311 --- /dev/null +++ b/test/test_data/query_auth/query_auth_split-null_rules Binary files differ