|
| 1 | +// Copyright 2026 Query Farm LLC - https://query.farm |
| 2 | + |
| 3 | +package farm.query.vgi; |
| 4 | + |
| 5 | +import farm.query.vgirpc.wire.Allocators; |
| 6 | +import org.apache.arrow.vector.BigIntVector; |
| 7 | +import org.apache.arrow.vector.BitVector; |
| 8 | +import org.apache.arrow.vector.FieldVector; |
| 9 | +import org.apache.arrow.vector.IntVector; |
| 10 | +import org.apache.arrow.vector.VarCharVector; |
| 11 | +import org.apache.arrow.vector.VectorSchemaRoot; |
| 12 | +import org.apache.arrow.vector.complex.StructVector; |
| 13 | +import org.apache.arrow.vector.ipc.ArrowStreamReader; |
| 14 | +import org.apache.arrow.vector.types.pojo.Field; |
| 15 | + |
| 16 | +import java.io.ByteArrayInputStream; |
| 17 | +import java.util.LinkedHashMap; |
| 18 | +import java.util.List; |
| 19 | +import java.util.Map; |
| 20 | +import java.util.Optional; |
| 21 | + |
| 22 | +/** |
| 23 | + * Resolved secrets passed to a worker, keyed by each secret's unique DuckDB |
| 24 | + * secret name (not by type) so several secrets of the same type (e.g. one per S3 |
| 25 | + * bucket) coexist. Each secret carries its connector-serialized {@code type} (the |
| 26 | + * DuckDB secret type) and {@code scope} (newline-joined scope prefixes) fields, |
| 27 | + * plus type-specific fields like {@code key_id}. |
| 28 | + * |
| 29 | + * <p>Mirrors {@code vgi::Secrets} in the Rust SDK. Parse the {@code byte[]} blob |
| 30 | + * carried on the params with {@link #parse(byte[])}, then select by name, type, |
| 31 | + * or scope.</p> |
| 32 | + */ |
| 33 | +public final class Secrets { |
| 34 | + |
| 35 | + /** name -> { field -> value-as-string }. */ |
| 36 | + private final Map<String, Map<String, String>> byName; |
| 37 | + |
| 38 | + private Secrets(Map<String, Map<String, String>> byName) { |
| 39 | + this.byName = byName; |
| 40 | + } |
| 41 | + |
| 42 | + /** |
| 43 | + * Build directly from a name -> fields map (for tests / non-IPC callers). |
| 44 | + * |
| 45 | + * @param byName the resolved secrets keyed by secret name |
| 46 | + * @return a {@code Secrets} over a copy of {@code byName} |
| 47 | + */ |
| 48 | + public static Secrets of(Map<String, Map<String, String>> byName) { |
| 49 | + Map<String, Map<String, String>> copy = new LinkedHashMap<>(); |
| 50 | + byName.forEach((k, v) -> copy.put(k, new LinkedHashMap<>(v))); |
| 51 | + return new Secrets(copy); |
| 52 | + } |
| 53 | + |
| 54 | + /** |
| 55 | + * Parse the IPC secrets blob. Each column is a secret (named by its DuckDB |
| 56 | + * secret name) holding a struct of its fields, including {@code type} and |
| 57 | + * {@code scope}. Empty/null blob yields empty secrets. |
| 58 | + * |
| 59 | + * @param bytes the IPC batch, or {@code null}/empty |
| 60 | + * @return the parsed secrets |
| 61 | + */ |
| 62 | + public static Secrets parse(byte[] bytes) { |
| 63 | + Map<String, Map<String, String>> byName = new LinkedHashMap<>(); |
| 64 | + if (bytes == null || bytes.length == 0) { |
| 65 | + return new Secrets(byName); |
| 66 | + } |
| 67 | + try (ByteArrayInputStream in = new ByteArrayInputStream(bytes); |
| 68 | + ArrowStreamReader r = new ArrowStreamReader(in, Allocators.root())) { |
| 69 | + if (r.loadNextBatch()) { |
| 70 | + VectorSchemaRoot root = r.getVectorSchemaRoot(); |
| 71 | + for (Field f : root.getSchema().getFields()) { |
| 72 | + FieldVector col = root.getVector(f.getName()); |
| 73 | + if (col == null || col.isNull(0)) { |
| 74 | + continue; |
| 75 | + } |
| 76 | + Map<String, String> fields = new LinkedHashMap<>(); |
| 77 | + if (col instanceof StructVector sv) { |
| 78 | + for (Field child : f.getChildren()) { |
| 79 | + FieldVector cv = sv.getChild(child.getName()); |
| 80 | + if (cv != null && !cv.isNull(0)) { |
| 81 | + fields.put(child.getName(), render(cv)); |
| 82 | + } |
| 83 | + } |
| 84 | + } else { |
| 85 | + fields.put(f.getName(), render(col)); |
| 86 | + } |
| 87 | + if (!fields.isEmpty()) { |
| 88 | + byName.put(f.getName(), fields); |
| 89 | + } |
| 90 | + } |
| 91 | + } |
| 92 | + } catch (Exception e) { |
| 93 | + throw new RuntimeException("Failed to parse secrets IPC batch", e); |
| 94 | + } |
| 95 | + return new Secrets(byName); |
| 96 | + } |
| 97 | + |
| 98 | + /** A field value from the first secret carrying it (any name). */ |
| 99 | + public Optional<String> field(String field) { |
| 100 | + return byName.values().stream() |
| 101 | + .map(m -> m.get(field)) |
| 102 | + .filter(java.util.Objects::nonNull) |
| 103 | + .findFirst(); |
| 104 | + } |
| 105 | + |
| 106 | + /** A named secret's field. */ |
| 107 | + public Optional<String> namedField(String name, String field) { |
| 108 | + Map<String, String> m = byName.get(name); |
| 109 | + return m == null ? Optional.empty() : Optional.ofNullable(m.get(field)); |
| 110 | + } |
| 111 | + |
| 112 | + /** Every resolved secret as (name -> fields). */ |
| 113 | + public Map<String, Map<String, String>> byName() { |
| 114 | + return byName; |
| 115 | + } |
| 116 | + |
| 117 | + /** The DuckDB secret type of the named secret (its {@code type} field). */ |
| 118 | + public Optional<String> secretType(String name) { |
| 119 | + return namedField(name, "type"); |
| 120 | + } |
| 121 | + |
| 122 | + /** Every resolved secret whose {@code type} field matches {@code secretType}. */ |
| 123 | + public List<Map<String, String>> ofType(String secretType) { |
| 124 | + return byName.values().stream() |
| 125 | + .filter(m -> secretType.equals(m.get("type"))) |
| 126 | + .toList(); |
| 127 | + } |
| 128 | + |
| 129 | + /** |
| 130 | + * The fields of the secret whose {@code scope} is the longest prefix of |
| 131 | + * {@code path}. The connector serializes each secret's scope as a |
| 132 | + * newline-joined list of prefixes; a secret with no (or empty) scope matches |
| 133 | + * as a last-resort fallback. Empty only when there are no candidate secrets. |
| 134 | + * |
| 135 | + * @param path the path to match (e.g. {@code s3://bucket/data/x.dat}) |
| 136 | + * @return the best-matching secret's fields |
| 137 | + */ |
| 138 | + public Optional<Map<String, String>> forScope(String path) { |
| 139 | + return selectForScope(path, null); |
| 140 | + } |
| 141 | + |
| 142 | + /** Like {@link #forScope} but only over secrets of {@code secretType}. */ |
| 143 | + public Optional<Map<String, String>> forScopeOfType(String path, String secretType) { |
| 144 | + return selectForScope(path, secretType); |
| 145 | + } |
| 146 | + |
| 147 | + /** A field of the best scope-matching secret for {@code path}. */ |
| 148 | + public Optional<String> fieldFor(String path, String field) { |
| 149 | + return forScope(path).map(m -> m.get(field)).filter(java.util.Objects::nonNull); |
| 150 | + } |
| 151 | + |
| 152 | + private Optional<Map<String, String>> selectForScope(String path, String secretType) { |
| 153 | + Map<String, String> best = null; |
| 154 | + int bestLen = -1; |
| 155 | + Map<String, String> fallback = null; |
| 156 | + for (Map<String, String> fields : byName.values()) { |
| 157 | + if (secretType != null && !secretType.equals(fields.get("type"))) { |
| 158 | + continue; |
| 159 | + } |
| 160 | + String scope = fields.get("scope"); |
| 161 | + if (scope == null || scope.isEmpty()) { |
| 162 | + if (fallback == null) { |
| 163 | + fallback = fields; |
| 164 | + } |
| 165 | + continue; |
| 166 | + } |
| 167 | + for (String prefix : scope.split("\n")) { |
| 168 | + if (!prefix.isEmpty() && path.startsWith(prefix) && prefix.length() > bestLen) { |
| 169 | + bestLen = prefix.length(); |
| 170 | + best = fields; |
| 171 | + } |
| 172 | + } |
| 173 | + } |
| 174 | + return Optional.ofNullable(best != null ? best : fallback); |
| 175 | + } |
| 176 | + |
| 177 | + private static String render(FieldVector v) { |
| 178 | + if (v.isNull(0)) { |
| 179 | + return ""; |
| 180 | + } |
| 181 | + if (v instanceof VarCharVector vc) { |
| 182 | + return new String(vc.get(0), java.nio.charset.StandardCharsets.UTF_8); |
| 183 | + } |
| 184 | + if (v instanceof BigIntVector iv) { |
| 185 | + return String.valueOf(iv.get(0)); |
| 186 | + } |
| 187 | + if (v instanceof IntVector iv) { |
| 188 | + return String.valueOf(iv.get(0)); |
| 189 | + } |
| 190 | + if (v instanceof BitVector bv) { |
| 191 | + return String.valueOf(bv.get(0) != 0); |
| 192 | + } |
| 193 | + Object o = v.getObject(0); |
| 194 | + return o == null ? "" : o.toString(); |
| 195 | + } |
| 196 | +} |
0 commit comments