|
| 1 | +// Copyright 2023 Google LLC |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +import * as protos from '../../protos/protos'; |
| 16 | +import {bqTypeToFieldTypeMap, convertModeToLabel} from './proto_mappings'; |
| 17 | + |
| 18 | +type TableSchema = protos.google.cloud.bigquery.storage.v1.ITableSchema; |
| 19 | +type TableFieldSchema = |
| 20 | + protos.google.cloud.bigquery.storage.v1.ITableFieldSchema; |
| 21 | +type FieldDescriptorProto = protos.google.protobuf.IFieldDescriptorProto; |
| 22 | +type FileDescriptorProto = protos.google.protobuf.IFileDescriptorProto; |
| 23 | +type FileDescriptorSet = protos.google.protobuf.FileDescriptorSet; |
| 24 | +type DescriptorProto = protos.google.protobuf.DescriptorProto; |
| 25 | +type FieldDescriptorProtoType = |
| 26 | + protos.google.protobuf.IFieldDescriptorProto['type']; |
| 27 | +type FieldDescriptorProtoLabel = |
| 28 | + protos.google.protobuf.IFieldDescriptorProto['label']; |
| 29 | + |
| 30 | +const TableFieldSchema = |
| 31 | + protos.google.cloud.bigquery.storage.v1.TableFieldSchema; |
| 32 | +const DescriptorProto = protos.google.protobuf.DescriptorProto; |
| 33 | +const FieldDescriptorProto = protos.google.protobuf.FieldDescriptorProto; |
| 34 | +const FileDescriptorProto = protos.google.protobuf.FileDescriptorProto; |
| 35 | +const FileDescriptorSet = protos.google.protobuf.FileDescriptorSet; |
| 36 | + |
| 37 | +const packedTypes: FieldDescriptorProtoType[] = [ |
| 38 | + FieldDescriptorProto.Type.TYPE_INT32, |
| 39 | + FieldDescriptorProto.Type.TYPE_INT64, |
| 40 | + FieldDescriptorProto.Type.TYPE_UINT32, |
| 41 | + FieldDescriptorProto.Type.TYPE_UINT64, |
| 42 | + FieldDescriptorProto.Type.TYPE_SINT32, |
| 43 | + FieldDescriptorProto.Type.TYPE_SINT64, |
| 44 | + FieldDescriptorProto.Type.TYPE_FIXED32, |
| 45 | + FieldDescriptorProto.Type.TYPE_FIXED64, |
| 46 | + FieldDescriptorProto.Type.TYPE_SFIXED32, |
| 47 | + FieldDescriptorProto.Type.TYPE_SFIXED64, |
| 48 | + FieldDescriptorProto.Type.TYPE_FLOAT, |
| 49 | + FieldDescriptorProto.Type.TYPE_DOUBLE, |
| 50 | + FieldDescriptorProto.Type.TYPE_BOOL, |
| 51 | + FieldDescriptorProto.Type.TYPE_ENUM, |
| 52 | +]; |
| 53 | + |
| 54 | +/** Builds a DescriptorProto for a given table schema using proto2 syntax. |
| 55 | + * @param schema - a BigQuery Storage TableSchema. |
| 56 | + * @param scope - scope to namespace protobuf structs. |
| 57 | + * @returns DescriptorProto |
| 58 | + */ |
| 59 | +export function convertStorageSchemaToProto2Descriptor( |
| 60 | + schema: TableSchema, |
| 61 | + scope: string |
| 62 | +): DescriptorProto { |
| 63 | + const fds = convertStorageSchemaToFileDescriptorInternal( |
| 64 | + schema, |
| 65 | + scope, |
| 66 | + false |
| 67 | + ); |
| 68 | + return normalizeDescriptorSet(fds); |
| 69 | +} |
| 70 | + |
| 71 | +/** Builds a DescriptorProto for a given table schema using proto3 syntax. |
| 72 | + * @param schema - a Bigquery TableSchema. |
| 73 | + * @param scope - scope to namespace protobuf structs. |
| 74 | + * @returns DescriptorProto |
| 75 | + */ |
| 76 | +export function convertStorageSchemaToProto3Descriptor( |
| 77 | + schema: TableSchema, |
| 78 | + scope: string |
| 79 | +): DescriptorProto { |
| 80 | + const fds = convertStorageSchemaToFileDescriptorInternal(schema, scope, true); |
| 81 | + return normalizeDescriptorSet(fds); |
| 82 | +} |
| 83 | + |
| 84 | +function convertStorageSchemaToFileDescriptorInternal( |
| 85 | + schema: TableSchema, |
| 86 | + scope: string, |
| 87 | + useProto3: boolean |
| 88 | +): FileDescriptorSet { |
| 89 | + let fNumber = 0; |
| 90 | + const fields: FieldDescriptorProto[] = []; |
| 91 | + const deps = new Map<string, FileDescriptorProto>(); |
| 92 | + for (const field of schema.fields ?? []) { |
| 93 | + fNumber += 1; |
| 94 | + const currentScope = `${scope}_${field.name}`; |
| 95 | + if (field.type === TableFieldSchema.Type.STRUCT) { |
| 96 | + const subSchema: TableSchema = { |
| 97 | + fields: field.fields, |
| 98 | + }; |
| 99 | + const fd = convertStorageSchemaToFileDescriptorInternal( |
| 100 | + subSchema, |
| 101 | + currentScope, |
| 102 | + useProto3 |
| 103 | + ); |
| 104 | + for (const f of fd.file) { |
| 105 | + if (f.name) { |
| 106 | + deps.set(f.name, f); |
| 107 | + } |
| 108 | + } |
| 109 | + const fdp = convertTableFieldSchemaToFieldDescriptorProto( |
| 110 | + field, |
| 111 | + fNumber, |
| 112 | + currentScope, |
| 113 | + useProto3 |
| 114 | + ); |
| 115 | + fields.push(fdp); |
| 116 | + } else { |
| 117 | + const fdp = convertTableFieldSchemaToFieldDescriptorProto( |
| 118 | + field, |
| 119 | + fNumber, |
| 120 | + currentScope, |
| 121 | + useProto3 |
| 122 | + ); |
| 123 | + fields.push(fdp); |
| 124 | + } |
| 125 | + } |
| 126 | + |
| 127 | + const dp = new DescriptorProto({ |
| 128 | + name: scope, |
| 129 | + field: fields, |
| 130 | + }); |
| 131 | + |
| 132 | + const depsNames: string[] = Array.from(deps.keys()); |
| 133 | + const syntax = useProto3 ? 'proto3' : 'proto2'; |
| 134 | + const fdp = new FileDescriptorProto({ |
| 135 | + messageType: [dp], |
| 136 | + name: `${scope}.proto`, |
| 137 | + syntax, |
| 138 | + dependency: depsNames, |
| 139 | + }); |
| 140 | + |
| 141 | + const fds = new FileDescriptorSet({ |
| 142 | + file: [fdp, ...Array.from(deps.values())], |
| 143 | + }); |
| 144 | + |
| 145 | + return fds; |
| 146 | +} |
| 147 | + |
| 148 | +function normalizeDescriptorSet(fds: FileDescriptorSet): DescriptorProto { |
| 149 | + let dp: DescriptorProto | null = null; |
| 150 | + let fdpName; |
| 151 | + if (fds.file.length > 0) { |
| 152 | + // search root descriptor |
| 153 | + const fdp = fds.file[0]; |
| 154 | + fdpName = fdp.name; |
| 155 | + if (fdp.messageType && fdp.messageType.length > 0) { |
| 156 | + dp = new DescriptorProto(fdp.messageType[0]); |
| 157 | + } |
| 158 | + } |
| 159 | + if (!dp) { |
| 160 | + throw Error('root descriptor not found'); |
| 161 | + } |
| 162 | + for (const fdp of fds.file) { |
| 163 | + if (fdp.name === fdpName) { |
| 164 | + continue; |
| 165 | + } |
| 166 | + if (!dp.nestedType) { |
| 167 | + dp.nestedType = []; |
| 168 | + } |
| 169 | + if (!fdp.messageType) { |
| 170 | + continue; |
| 171 | + } |
| 172 | + for (const nestedDP of fdp.messageType) { |
| 173 | + dp.nestedType.push(normalizeDescriptor(new DescriptorProto(nestedDP))); |
| 174 | + } |
| 175 | + } |
| 176 | + return normalizeDescriptor(dp); |
| 177 | +} |
| 178 | + |
| 179 | +/** |
| 180 | + * Builds a self-contained DescriptorProto suitable for communicating schema |
| 181 | + * information with the BigQuery Storage write API. It's primarily used for cases where users are |
| 182 | + * interested in sending data using a predefined protocol buffer message. |
| 183 | + * @param dp - DescriptorProto to be bundled. |
| 184 | + * @return DescriptorProto |
| 185 | + */ |
| 186 | +export function normalizeDescriptor(dp: DescriptorProto): DescriptorProto { |
| 187 | + dp.name = normalizeName(dp.name); |
| 188 | + for (const f of dp.field) { |
| 189 | + if (f.proto3Optional) { |
| 190 | + f.proto3Optional = null; |
| 191 | + } |
| 192 | + if (f.oneofIndex) { |
| 193 | + f.oneofIndex = null; |
| 194 | + } |
| 195 | + if (f.options) { |
| 196 | + f.options.packed = shouldPackType(f.type, f.label, false); |
| 197 | + } |
| 198 | + } |
| 199 | + const normalizedNestedTypes = []; |
| 200 | + for (const nestedDP of dp.nestedType) { |
| 201 | + normalizedNestedTypes.push( |
| 202 | + normalizeDescriptor(new DescriptorProto(nestedDP)) |
| 203 | + ); |
| 204 | + } |
| 205 | + dp.nestedType = normalizedNestedTypes; |
| 206 | + return dp; |
| 207 | +} |
| 208 | + |
| 209 | +function normalizeName(name: string): string { |
| 210 | + return name.replace(/\./, '_'); |
| 211 | +} |
| 212 | + |
| 213 | +function convertTableFieldSchemaToFieldDescriptorProto( |
| 214 | + field: TableFieldSchema, |
| 215 | + fNumber: number, |
| 216 | + scope: string, |
| 217 | + useProto3: boolean |
| 218 | +): FieldDescriptorProto { |
| 219 | + const name = `${field.name}`.toLowerCase(); |
| 220 | + const type = field.type; |
| 221 | + if (!type) { |
| 222 | + throw Error(`table field ${name} missing type`); |
| 223 | + } |
| 224 | + const label = convertModeToLabel(field.mode, useProto3); |
| 225 | + let fdp: FieldDescriptorProto; |
| 226 | + if (type === TableFieldSchema.Type.STRUCT) { |
| 227 | + fdp = new FieldDescriptorProto({ |
| 228 | + name: name, |
| 229 | + number: fNumber, |
| 230 | + type: FieldDescriptorProto.Type.TYPE_MESSAGE, |
| 231 | + typeName: scope, |
| 232 | + label: label, |
| 233 | + }); |
| 234 | + } else { |
| 235 | + const pType = bqTypeToFieldTypeMap[type]; |
| 236 | + if (pType === null) { |
| 237 | + throw Error(`table field type ${type} not supported`); |
| 238 | + } |
| 239 | + fdp = new FieldDescriptorProto({ |
| 240 | + name: field.name, |
| 241 | + number: fNumber, |
| 242 | + type: pType, |
| 243 | + label: label, |
| 244 | + options: { |
| 245 | + packed: shouldPackType(pType, label, useProto3), |
| 246 | + }, |
| 247 | + proto3Optional: isProto3Optional(label, useProto3), |
| 248 | + }); |
| 249 | + } |
| 250 | + return fdp; |
| 251 | +} |
| 252 | + |
| 253 | +function shouldPackType( |
| 254 | + t: FieldDescriptorProtoType, |
| 255 | + label: FieldDescriptorProtoLabel | null, |
| 256 | + useProto3: boolean |
| 257 | +): boolean | undefined { |
| 258 | + if (useProto3) { |
| 259 | + return false; |
| 260 | + } |
| 261 | + if (label !== FieldDescriptorProto.Label.LABEL_REPEATED) { |
| 262 | + return undefined; |
| 263 | + } |
| 264 | + return packedTypes.includes(t); |
| 265 | +} |
| 266 | + |
| 267 | +function isProto3Optional( |
| 268 | + label: FieldDescriptorProtoLabel | null, |
| 269 | + useProto3: boolean |
| 270 | +): boolean | undefined { |
| 271 | + if (!useProto3) { |
| 272 | + return undefined; |
| 273 | + } |
| 274 | + return label === FieldDescriptorProto.Label.LABEL_OPTIONAL; |
| 275 | +} |
0 commit comments