(WIP) adjust routines logic, fix tests

Author: plamut

google/cloud/bigquery/routine/routine.py

@@ -16,11 +16,10 @@
 
 """Define resources for the BigQuery Routines API."""
 
-from google.protobuf import json_format
-
 import google.cloud._helpers
 from google.cloud.bigquery import _helpers
 from google.cloud.bigquery.types import StandardSqlDataType
+from google.cloud.bigquery.types import StandardSqlField
 from google.cloud.bigquery.types import StandardSqlTableType
 
 
@@ -206,16 +205,12 @@ def return_type(self):
         if not resource:
             return resource
 
-        output = StandardSqlDataType()
-        raw_protobuf = json_format.ParseDict(
-            resource, output._pb, ignore_unknown_fields=True
-        )
-        return type(output).wrap(raw_protobuf)
+        return StandardSqlDataType.from_api_repr(resource)
 
     @return_type.setter
-    def return_type(self, value):
+    def return_type(self, value: StandardSqlDataType):
         if value:
-            resource = json_format.MessageToDict(value._pb)
+            resource = value.to_api_repr()
         else:
             resource = None
         self._properties[self._PROPERTY_TO_API_FIELD["return_type"]] = resource
@@ -232,20 +227,36 @@ def return_table_type(self) -> StandardSqlTableType:
         if not resource:
             return resource
 
-        output = StandardSqlTableType()
-        raw_protobuf = json_format.ParseDict(
-            resource, output._pb, ignore_unknown_fields=True
+        # TODO: this should be StandardSqlTableType.from_api_repr()
+        columns = (
+            StandardSqlField(
+                name=column.get("name"),
+                type=StandardSqlDataType.from_api_repr(column["type"]),
+            )
+            for column in resource.get("columns")
         )
-        return type(output).wrap(raw_protobuf)
+        # output = StandardSqlTableType()
+        # raw_protobuf = json_format.ParseDict(
+        #     resource, output._pb, ignore_unknown_fields=True
+        # )
+        # return type(output).wrap(raw_protobuf)
+        return StandardSqlTableType(columns=columns)
 
     @return_table_type.setter
     def return_table_type(self, value):
         if not value:
             resource = None
         else:
-            resource = {
-                "columns": [json_format.MessageToDict(col._pb) for col in value.columns]
-            }
+            columns = [
+                # TODO: this should really be StandardSqlField.to_api_repr()
+                # ACTUALLY, all this thould be wrapped in StandardSqlTableType.to_api_repr()
+                {"name": col.name, "type": col.type.to_api_repr()}
+                for col in value.columns
+            ]
+            # resource = {
+            #     "columns": [json_format.MessageToDict(col._pb) for col in value.columns]
+            # }
+            resource = {"columns": columns}
 
         self._properties[self._PROPERTY_TO_API_FIELD["return_table_type"]] = resource
 
@@ -417,16 +428,17 @@ def data_type(self):
         if not resource:
             return resource
 
-        output = StandardSqlDataType()
-        raw_protobuf = json_format.ParseDict(
-            resource, output._pb, ignore_unknown_fields=True
-        )
-        return type(output).wrap(raw_protobuf)
+        # output = StandardSqlDataType()
+        # raw_protobuf = json_format.ParseDict(
+        #     resource, output._pb, ignore_unknown_fields=True
+        # )
+        # return type(output).wrap(raw_protobuf)
+        return StandardSqlDataType.from_api_repr(resource)
 
     @data_type.setter
     def data_type(self, value):
         if value:
-            resource = json_format.MessageToDict(value._pb)
+            resource = value.to_api_repr()
         else:
             resource = None
         self._properties[self._PROPERTY_TO_API_FIELD["data_type"]] = resource

google/cloud/bigquery/types/standard_sql.py

@@ -78,6 +78,63 @@ def __post_init__(self):
                 "array_element_type and struct_type are mutally exclusive."
             )
 
+    def to_api_repr(self):
+        """TODO: docstirng... and tests
+
+        mention link in class docstring?  https://cloud.google.com/bigqPuery/docs/reference/rest/v2/StandardSqlDataType
+        """
+        if not self.type_kind:
+            type_kind = self.TypeKind.TYPE_KIND_UNSPECIFIED.value
+        else:
+            type_kind = self.type_kind.value
+
+        result = {"typeKind": type_kind}
+
+        if self.type_kind == self.TypeKind.ARRAY:
+            if not self.array_element_type:
+                array_type = None
+            else:
+                array_type = self.array_element_type.to_api_repr()
+            result["arrayElementType"] = array_type
+        elif self.type_kind == self.TypeKind.STRUCT:
+            if not self.struct_type:
+                struct_type = None
+            else:
+                struct_type = self.struct_type.to_api_repr()
+            result["structType"] = struct_type
+
+        return result
+
+    @classmethod
+    def from_api_repr(cls, resource):
+        """TODO: type annotations, tests..."""
+        type_kind = resource.get("typeKind")
+        if type_kind not in cls.TypeKind.__members__:
+            type_kind = cls.TypeKind.TYPE_KIND_UNSPECIFIED
+        else:
+            type_kind = cls.TypeKind[type_kind]  # From string to enum member.
+
+        array_element_type = None
+        if type_kind == cls.type_kind.ARRAY:
+            array_element_type = cls.from_api_repr(resource["arrayElementType"])
+
+        struct_type = None
+        if type_kind == cls.TypeKind.STRUCT:
+            # TODO: this logic belongs to StandardSqlStructType.to_api_repr()
+            struct_info = resource["structType"]
+            fields = struct_info.get("fields", [])
+
+            fields = (
+                StandardSqlField(
+                    name=field.get("name"),
+                    type=StandardSqlDataType.from_api_repr(field.get["type"]),
+                )
+                for field in struct_info.get("fields", [])
+            )
+            struct_type = StandardSqlStructType(fields=fields)
+
+        return cls(type_kind, array_element_type, struct_type)
+
 
 @dataclass
 class StandardSqlField: