|
1 | 1 | # Copyright 2019-present Kensho Technologies, LLC. |
2 | | -from typing import Dict, NamedTuple, Set |
| 2 | +from dataclasses import dataclass |
| 3 | +from typing import AbstractSet, Dict, Set, Tuple, Union |
3 | 4 | import warnings |
4 | 5 |
|
5 | 6 | import six |
6 | 7 | from sqlalchemy import Table |
7 | 8 |
|
8 | 9 | from ...schema import INBOUND_EDGE_FIELD_PREFIX, OUTBOUND_EDGE_FIELD_PREFIX |
9 | | -from ...schema.schema_info import DirectJoinDescriptor |
| 10 | +from ...schema.schema_info import CompositeJoinDescriptor, DirectJoinDescriptor, JoinDescriptor |
10 | 11 | from ..exceptions import InvalidSQLEdgeError |
11 | 12 |
|
12 | 13 |
|
13 | | -class DirectEdgeDescriptor(NamedTuple): |
| 14 | +@dataclass(frozen=True) |
| 15 | +class DirectEdgeDescriptor: |
| 16 | + """Represents a bidirectional edge between two vertices backed by DirectJoinDescriptors.""" |
| 17 | + |
14 | 18 | from_vertex: str # Name of the source vertex. |
15 | 19 | from_column: str # Name of the column of the underlying source table to use for SQL join. |
16 | 20 | to_vertex: str # Name of the destination vertex. |
17 | 21 | to_column: str # Name of the column of the underlying destination table to use for SQL join. |
18 | 22 |
|
19 | 23 |
|
| 24 | +@dataclass(frozen=True) |
| 25 | +class CompositeEdgeDescriptor: |
| 26 | + """Represents a bidirectional edge between two vertices backed by CompositeJoinDescriptors.""" |
| 27 | + |
| 28 | + from_vertex: str # Name of the source vertex |
| 29 | + to_vertex: str # Name of the destination vertex |
| 30 | + |
| 31 | + # (from_column, to_column) pairs, where from_column is on the origin table |
| 32 | + # and to_column is on the destination table of the join. |
| 33 | + column_pairs: AbstractSet[Tuple[str, str]] |
| 34 | + |
| 35 | + def __post_init__(self) -> None: |
| 36 | + """Validate fields.""" |
| 37 | + if not self.column_pairs: |
| 38 | + raise AssertionError("The column_pairs field is expected to be non-empty.") |
| 39 | + |
| 40 | + |
| 41 | +EdgeDescriptor = Union[DirectEdgeDescriptor, CompositeEdgeDescriptor] |
| 42 | + |
| 43 | + |
20 | 44 | def get_join_descriptors_from_edge_descriptors( |
21 | | - direct_edges: Dict[str, DirectEdgeDescriptor] |
22 | | -) -> Dict[str, Dict[str, DirectJoinDescriptor]]: |
| 45 | + direct_edges: Dict[str, EdgeDescriptor] |
| 46 | +) -> Dict[str, Dict[str, JoinDescriptor]]: |
23 | 47 | """Return the SQL edges in a format more suited to resolving vertex fields.""" |
24 | | - join_descriptors: Dict[str, Dict[str, DirectJoinDescriptor]] = {} |
25 | | - for edge_name, direct_edge_descriptor in direct_edges.items(): |
26 | | - from_column = direct_edge_descriptor.from_column |
27 | | - to_column = direct_edge_descriptor.to_column |
28 | | - join_descriptors.setdefault(direct_edge_descriptor.from_vertex, {}) |
29 | | - join_descriptors.setdefault(direct_edge_descriptor.to_vertex, {}) |
| 48 | + join_descriptors: Dict[str, Dict[str, JoinDescriptor]] = {} |
| 49 | + for edge_name, edge_descriptor in direct_edges.items(): |
| 50 | + join_descriptors.setdefault(edge_descriptor.from_vertex, {}) |
| 51 | + join_descriptors.setdefault(edge_descriptor.to_vertex, {}) |
30 | 52 | out_edge_name = OUTBOUND_EDGE_FIELD_PREFIX + edge_name |
31 | 53 | in_edge_name = INBOUND_EDGE_FIELD_PREFIX + edge_name |
32 | | - join_descriptors[direct_edge_descriptor.from_vertex][out_edge_name] = DirectJoinDescriptor( |
33 | | - from_column, to_column |
34 | | - ) |
35 | | - join_descriptors[direct_edge_descriptor.to_vertex][in_edge_name] = DirectJoinDescriptor( |
36 | | - to_column, from_column |
37 | | - ) |
| 54 | + if isinstance(edge_descriptor, DirectEdgeDescriptor): |
| 55 | + from_column = edge_descriptor.from_column |
| 56 | + to_column = edge_descriptor.to_column |
| 57 | + join_descriptors[edge_descriptor.from_vertex][out_edge_name] = DirectJoinDescriptor( |
| 58 | + from_column, to_column |
| 59 | + ) |
| 60 | + join_descriptors[edge_descriptor.to_vertex][in_edge_name] = DirectJoinDescriptor( |
| 61 | + to_column, from_column |
| 62 | + ) |
| 63 | + elif isinstance(edge_descriptor, CompositeEdgeDescriptor): |
| 64 | + join_descriptors[edge_descriptor.from_vertex][out_edge_name] = CompositeJoinDescriptor( |
| 65 | + edge_descriptor.column_pairs |
| 66 | + ) |
| 67 | + join_descriptors[edge_descriptor.to_vertex][in_edge_name] = CompositeJoinDescriptor( |
| 68 | + { |
| 69 | + (to_column, from_column) |
| 70 | + for from_column, to_column in edge_descriptor.column_pairs |
| 71 | + } |
| 72 | + ) |
| 73 | + else: |
| 74 | + raise AssertionError( |
| 75 | + f"Unknown edge descriptor type {edge_descriptor}: " |
| 76 | + f"{type(edge_descriptor)} for edge {edge_name}." |
| 77 | + ) |
38 | 78 | return join_descriptors |
39 | 79 |
|
40 | 80 |
|
41 | 81 | def validate_edge_descriptors( |
42 | | - vertex_name_to_table: Dict[str, Table], direct_edges: Dict[str, DirectEdgeDescriptor] |
| 82 | + vertex_name_to_table: Dict[str, Table], edges: Dict[str, EdgeDescriptor] |
43 | 83 | ) -> None: |
44 | 84 | """Validate that the edge descriptors do not reference non-existent vertices or columns.""" |
45 | 85 | # TODO(pmantica1): Validate that columns in a direct SQL edge have comparable types. |
46 | 86 | # TODO(pmantica1): Validate that columns don't have types that probably shouldn't be used for |
47 | 87 | # joins, (e.g. array types). |
48 | | - for edge_name, direct_edge_descriptor in six.iteritems(direct_edges): |
49 | | - for vertex_name, column_name in ( |
50 | | - (direct_edge_descriptor.from_vertex, direct_edge_descriptor.from_column), |
51 | | - (direct_edge_descriptor.to_vertex, direct_edge_descriptor.to_column), |
52 | | - ): |
| 88 | + for edge_name, edge_descriptor in six.iteritems(edges): |
| 89 | + if isinstance(edge_descriptor, DirectEdgeDescriptor): |
| 90 | + vertex_column_pairs = [ |
| 91 | + (edge_descriptor.from_vertex, edge_descriptor.from_column), |
| 92 | + (edge_descriptor.to_vertex, edge_descriptor.to_column), |
| 93 | + ] |
| 94 | + elif isinstance(edge_descriptor, CompositeEdgeDescriptor): |
| 95 | + vertex_column_pairs = [ |
| 96 | + (edge_descriptor.from_vertex, from_column) |
| 97 | + for from_column, _ in edge_descriptor.column_pairs |
| 98 | + ] + [ |
| 99 | + (edge_descriptor.to_vertex, to_column) |
| 100 | + for _, to_column in edge_descriptor.column_pairs |
| 101 | + ] |
| 102 | + else: |
| 103 | + raise AssertionError( |
| 104 | + f"Unknown edge descriptor type {edge_descriptor}: " |
| 105 | + f"{type(edge_descriptor)} for edge {edge_name}." |
| 106 | + ) |
| 107 | + |
| 108 | + for vertex_name, column_name in vertex_column_pairs: |
53 | 109 | if vertex_name not in vertex_name_to_table: |
54 | 110 | raise InvalidSQLEdgeError( |
55 | 111 | "SQL edge {} with edge descriptor {} references a " |
56 | | - "non-existent vertex {}".format(edge_name, direct_edge_descriptor, vertex_name) |
| 112 | + "non-existent vertex {}".format(edge_name, edge_descriptor, vertex_name) |
57 | 113 | ) |
58 | 114 | if column_name not in vertex_name_to_table[vertex_name].columns: |
59 | 115 | raise InvalidSQLEdgeError( |
60 | 116 | "SQL edge {} with edge descriptor {} references a " |
61 | | - "non-existent column {}".format(edge_name, direct_edge_descriptor, column_name) |
| 117 | + "non-existent column {}".format(edge_name, edge_descriptor, column_name) |
62 | 118 | ) |
63 | 119 |
|
64 | 120 |
|
|
0 commit comments