390 lines
13 KiB
Python
390 lines
13 KiB
Python
# Copyright 2021 Google LLC
|
|
|
|
# Licensed 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
|
|
|
|
# https://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 copy
|
|
import typing
|
|
from typing import Any, Dict, Iterable, List, Optional
|
|
|
|
from google.cloud.bigquery.enums import StandardSqlTypeNames
|
|
|
|
|
|
class StandardSqlDataType:
|
|
"""The type of a variable, e.g., a function argument.
|
|
|
|
See:
|
|
https://cloud.google.com/bigquery/docs/reference/rest/v2/StandardSqlDataType
|
|
|
|
Examples:
|
|
|
|
.. code-block:: text
|
|
|
|
INT64: {type_kind="INT64"}
|
|
ARRAY: {type_kind="ARRAY", array_element_type="STRING"}
|
|
STRUCT<x STRING, y ARRAY>: {
|
|
type_kind="STRUCT",
|
|
struct_type={
|
|
fields=[
|
|
{name="x", type={type_kind="STRING"}},
|
|
{
|
|
name="y",
|
|
type={type_kind="ARRAY", array_element_type="DATE"}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
RANGE: {type_kind="RANGE", range_element_type="DATETIME"}
|
|
|
|
Args:
|
|
type_kind:
|
|
The top level type of this field. Can be any standard SQL data type,
|
|
e.g. INT64, DATE, ARRAY.
|
|
array_element_type:
|
|
The type of the array's elements, if type_kind is ARRAY.
|
|
struct_type:
|
|
The fields of this struct, in order, if type_kind is STRUCT.
|
|
range_element_type:
|
|
The type of the range's elements, if type_kind is RANGE.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
type_kind: Optional[
|
|
StandardSqlTypeNames
|
|
] = StandardSqlTypeNames.TYPE_KIND_UNSPECIFIED,
|
|
array_element_type: Optional["StandardSqlDataType"] = None,
|
|
struct_type: Optional["StandardSqlStructType"] = None,
|
|
range_element_type: Optional["StandardSqlDataType"] = None,
|
|
):
|
|
self._properties: Dict[str, Any] = {}
|
|
|
|
self.type_kind = type_kind
|
|
self.array_element_type = array_element_type
|
|
self.struct_type = struct_type
|
|
self.range_element_type = range_element_type
|
|
|
|
@property
|
|
def type_kind(self) -> Optional[StandardSqlTypeNames]:
|
|
"""The top level type of this field.
|
|
|
|
Can be any standard SQL data type, e.g. INT64, DATE, ARRAY.
|
|
"""
|
|
kind = self._properties["typeKind"]
|
|
return StandardSqlTypeNames[kind] # pytype: disable=missing-parameter
|
|
|
|
@type_kind.setter
|
|
def type_kind(self, value: Optional[StandardSqlTypeNames]):
|
|
if not value:
|
|
kind = StandardSqlTypeNames.TYPE_KIND_UNSPECIFIED.value
|
|
else:
|
|
kind = value.value
|
|
self._properties["typeKind"] = kind
|
|
|
|
@property
|
|
def array_element_type(self) -> Optional["StandardSqlDataType"]:
|
|
"""The type of the array's elements, if type_kind is ARRAY."""
|
|
element_type = self._properties.get("arrayElementType")
|
|
|
|
if element_type is None:
|
|
return None
|
|
|
|
result = StandardSqlDataType()
|
|
result._properties = element_type # We do not use a copy on purpose.
|
|
return result
|
|
|
|
@array_element_type.setter
|
|
def array_element_type(self, value: Optional["StandardSqlDataType"]):
|
|
element_type = None if value is None else value.to_api_repr()
|
|
|
|
if element_type is None:
|
|
self._properties.pop("arrayElementType", None)
|
|
else:
|
|
self._properties["arrayElementType"] = element_type
|
|
|
|
@property
|
|
def struct_type(self) -> Optional["StandardSqlStructType"]:
|
|
"""The fields of this struct, in order, if type_kind is STRUCT."""
|
|
struct_info = self._properties.get("structType")
|
|
|
|
if struct_info is None:
|
|
return None
|
|
|
|
result = StandardSqlStructType()
|
|
result._properties = struct_info # We do not use a copy on purpose.
|
|
return result
|
|
|
|
@struct_type.setter
|
|
def struct_type(self, value: Optional["StandardSqlStructType"]):
|
|
struct_type = None if value is None else value.to_api_repr()
|
|
|
|
if struct_type is None:
|
|
self._properties.pop("structType", None)
|
|
else:
|
|
self._properties["structType"] = struct_type
|
|
|
|
@property
|
|
def range_element_type(self) -> Optional["StandardSqlDataType"]:
|
|
"""The type of the range's elements, if type_kind = "RANGE". Must be
|
|
one of DATETIME, DATE, or TIMESTAMP."""
|
|
range_element_info = self._properties.get("rangeElementType")
|
|
|
|
if range_element_info is None:
|
|
return None
|
|
|
|
result = StandardSqlDataType()
|
|
result._properties = range_element_info # We do not use a copy on purpose.
|
|
return result
|
|
|
|
@range_element_type.setter
|
|
def range_element_type(self, value: Optional["StandardSqlDataType"]):
|
|
range_element_type = None if value is None else value.to_api_repr()
|
|
|
|
if range_element_type is None:
|
|
self._properties.pop("rangeElementType", None)
|
|
else:
|
|
self._properties["rangeElementType"] = range_element_type
|
|
|
|
def to_api_repr(self) -> Dict[str, Any]:
|
|
"""Construct the API resource representation of this SQL data type."""
|
|
return copy.deepcopy(self._properties)
|
|
|
|
@classmethod
|
|
def from_api_repr(cls, resource: Dict[str, Any]):
|
|
"""Construct an SQL data type instance given its API representation."""
|
|
type_kind = resource.get("typeKind")
|
|
if type_kind not in StandardSqlTypeNames.__members__:
|
|
type_kind = StandardSqlTypeNames.TYPE_KIND_UNSPECIFIED
|
|
else:
|
|
# Convert string to an enum member.
|
|
type_kind = StandardSqlTypeNames[ # pytype: disable=missing-parameter
|
|
typing.cast(str, type_kind)
|
|
]
|
|
|
|
array_element_type = None
|
|
if type_kind == StandardSqlTypeNames.ARRAY:
|
|
element_type = resource.get("arrayElementType")
|
|
if element_type:
|
|
array_element_type = cls.from_api_repr(element_type)
|
|
|
|
struct_type = None
|
|
if type_kind == StandardSqlTypeNames.STRUCT:
|
|
struct_info = resource.get("structType")
|
|
if struct_info:
|
|
struct_type = StandardSqlStructType.from_api_repr(struct_info)
|
|
|
|
range_element_type = None
|
|
if type_kind == StandardSqlTypeNames.RANGE:
|
|
range_element_info = resource.get("rangeElementType")
|
|
if range_element_info:
|
|
range_element_type = cls.from_api_repr(range_element_info)
|
|
|
|
return cls(type_kind, array_element_type, struct_type, range_element_type)
|
|
|
|
def __eq__(self, other):
|
|
if not isinstance(other, StandardSqlDataType):
|
|
return NotImplemented
|
|
else:
|
|
return (
|
|
self.type_kind == other.type_kind
|
|
and self.array_element_type == other.array_element_type
|
|
and self.struct_type == other.struct_type
|
|
and self.range_element_type == other.range_element_type
|
|
)
|
|
|
|
def __str__(self):
|
|
result = f"{self.__class__.__name__}(type_kind={self.type_kind!r}, ...)"
|
|
return result
|
|
|
|
|
|
class StandardSqlField:
|
|
"""A field or a column.
|
|
|
|
See:
|
|
https://cloud.google.com/bigquery/docs/reference/rest/v2/StandardSqlField
|
|
|
|
Args:
|
|
name:
|
|
The name of this field. Can be absent for struct fields.
|
|
type:
|
|
The type of this parameter. Absent if not explicitly specified.
|
|
|
|
For example, CREATE FUNCTION statement can omit the return type; in this
|
|
case the output parameter does not have this "type" field).
|
|
"""
|
|
|
|
def __init__(
|
|
self, name: Optional[str] = None, type: Optional[StandardSqlDataType] = None
|
|
):
|
|
type_repr = None if type is None else type.to_api_repr()
|
|
self._properties = {"name": name, "type": type_repr}
|
|
|
|
@property
|
|
def name(self) -> Optional[str]:
|
|
"""The name of this field. Can be absent for struct fields."""
|
|
return typing.cast(Optional[str], self._properties["name"])
|
|
|
|
@name.setter
|
|
def name(self, value: Optional[str]):
|
|
self._properties["name"] = value
|
|
|
|
@property
|
|
def type(self) -> Optional[StandardSqlDataType]:
|
|
"""The type of this parameter. Absent if not explicitly specified.
|
|
|
|
For example, CREATE FUNCTION statement can omit the return type; in this
|
|
case the output parameter does not have this "type" field).
|
|
"""
|
|
type_info = self._properties["type"]
|
|
|
|
if type_info is None:
|
|
return None
|
|
|
|
result = StandardSqlDataType()
|
|
# We do not use a properties copy on purpose.
|
|
result._properties = typing.cast(Dict[str, Any], type_info)
|
|
|
|
return result
|
|
|
|
@type.setter
|
|
def type(self, value: Optional[StandardSqlDataType]):
|
|
value_repr = None if value is None else value.to_api_repr()
|
|
self._properties["type"] = value_repr
|
|
|
|
def to_api_repr(self) -> Dict[str, Any]:
|
|
"""Construct the API resource representation of this SQL field."""
|
|
return copy.deepcopy(self._properties)
|
|
|
|
@classmethod
|
|
def from_api_repr(cls, resource: Dict[str, Any]):
|
|
"""Construct an SQL field instance given its API representation."""
|
|
result = cls(
|
|
name=resource.get("name"),
|
|
type=StandardSqlDataType.from_api_repr(resource.get("type", {})),
|
|
)
|
|
return result
|
|
|
|
def __eq__(self, other):
|
|
if not isinstance(other, StandardSqlField):
|
|
return NotImplemented
|
|
else:
|
|
return self.name == other.name and self.type == other.type
|
|
|
|
|
|
class StandardSqlStructType:
|
|
"""Type of a struct field.
|
|
|
|
See:
|
|
https://cloud.google.com/bigquery/docs/reference/rest/v2/StandardSqlDataType#StandardSqlStructType
|
|
|
|
Args:
|
|
fields: The fields in this struct.
|
|
"""
|
|
|
|
def __init__(self, fields: Optional[Iterable[StandardSqlField]] = None):
|
|
if fields is None:
|
|
fields = []
|
|
self._properties = {"fields": [field.to_api_repr() for field in fields]}
|
|
|
|
@property
|
|
def fields(self) -> List[StandardSqlField]:
|
|
"""The fields in this struct."""
|
|
result = []
|
|
|
|
for field_resource in self._properties.get("fields", []):
|
|
field = StandardSqlField()
|
|
field._properties = field_resource # We do not use a copy on purpose.
|
|
result.append(field)
|
|
|
|
return result
|
|
|
|
@fields.setter
|
|
def fields(self, value: Iterable[StandardSqlField]):
|
|
self._properties["fields"] = [field.to_api_repr() for field in value]
|
|
|
|
def to_api_repr(self) -> Dict[str, Any]:
|
|
"""Construct the API resource representation of this SQL struct type."""
|
|
return copy.deepcopy(self._properties)
|
|
|
|
@classmethod
|
|
def from_api_repr(cls, resource: Dict[str, Any]) -> "StandardSqlStructType":
|
|
"""Construct an SQL struct type instance given its API representation."""
|
|
fields = (
|
|
StandardSqlField.from_api_repr(field_resource)
|
|
for field_resource in resource.get("fields", [])
|
|
)
|
|
return cls(fields=fields)
|
|
|
|
def __eq__(self, other):
|
|
if not isinstance(other, StandardSqlStructType):
|
|
return NotImplemented
|
|
else:
|
|
return self.fields == other.fields
|
|
|
|
|
|
class StandardSqlTableType:
|
|
"""A table type.
|
|
|
|
See:
|
|
https://cloud.google.com/workflows/docs/reference/googleapis/bigquery/v2/Overview#StandardSqlTableType
|
|
|
|
Args:
|
|
columns: The columns in this table type.
|
|
"""
|
|
|
|
def __init__(self, columns: Iterable[StandardSqlField]):
|
|
self._properties = {"columns": [col.to_api_repr() for col in columns]}
|
|
|
|
@property
|
|
def columns(self) -> List[StandardSqlField]:
|
|
"""The columns in this table type."""
|
|
result = []
|
|
|
|
for column_resource in self._properties.get("columns", []):
|
|
column = StandardSqlField()
|
|
column._properties = column_resource # We do not use a copy on purpose.
|
|
result.append(column)
|
|
|
|
return result
|
|
|
|
@columns.setter
|
|
def columns(self, value: Iterable[StandardSqlField]):
|
|
self._properties["columns"] = [col.to_api_repr() for col in value]
|
|
|
|
def to_api_repr(self) -> Dict[str, Any]:
|
|
"""Construct the API resource representation of this SQL table type."""
|
|
return copy.deepcopy(self._properties)
|
|
|
|
@classmethod
|
|
def from_api_repr(cls, resource: Dict[str, Any]) -> "StandardSqlTableType":
|
|
"""Construct an SQL table type instance given its API representation."""
|
|
columns = []
|
|
|
|
for column_resource in resource.get("columns", []):
|
|
type_ = column_resource.get("type")
|
|
if type_ is None:
|
|
type_ = {}
|
|
|
|
column = StandardSqlField(
|
|
name=column_resource.get("name"),
|
|
type=StandardSqlDataType.from_api_repr(type_),
|
|
)
|
|
columns.append(column)
|
|
|
|
return cls(columns=columns)
|
|
|
|
def __eq__(self, other):
|
|
if not isinstance(other, StandardSqlTableType):
|
|
return NotImplemented
|
|
else:
|
|
return self.columns == other.columns
|