feat(spanner): add interval type support

Author: rahul2393

google/cloud/spanner_v1/_helpers.py

@@ -20,6 +20,8 @@
 import time
 import base64
 import threading
+import re
+from dataclasses import dataclass
 
 from google.protobuf.struct_pb2 import ListValue
 from google.protobuf.struct_pb2 import Value
@@ -196,6 +198,162 @@ def _datetime_to_rfc3339_nanoseconds(value):
     return "{}.{}Z".format(value.isoformat(sep="T", timespec="seconds"), nanos)
 
 
+@dataclass
+class Interval:
+    """Represents a Spanner INTERVAL type.
+    
+    An interval is a combination of months, days and nanoseconds.
+    Internally, Spanner supports Interval value with the following range of individual fields:
+    months: [-120000, 120000]
+    days: [-3660000, 3660000]
+    nanoseconds: [-316224000000000000000, 316224000000000000000]
+    """
+    months: int = 0
+    days: int = 0
+    nanos: int = 0
+
+    def __str__(self) -> str:
+        """Returns the ISO8601 duration format string representation."""
+        result = ["P"]
+        
+        # Handle years and months
+        if self.months:
+            is_negative = self.months < 0
+            abs_months = abs(self.months)
+            years, months = divmod(abs_months, 12)
+            if years:
+                result.append(f"{'-' if is_negative else ''}{years}Y")
+            if months:
+                result.append(f"{'-' if is_negative else ''}{months}M")
+            
+        # Handle days
+        if self.days:
+            result.append(f"{self.days}D")
+            
+        # Handle time components
+        if self.nanos:
+            result.append("T")
+            nanos = abs(self.nanos)
+            is_negative = self.nanos < 0
+            
+            # Convert to hours, minutes, seconds
+            nanos_per_hour = 3600000000000
+            hours, nanos = divmod(nanos, nanos_per_hour)
+            if hours:
+                if is_negative:
+                    result.append("-")
+                result.append(f"{hours}H")
+                
+            nanos_per_minute = 60000000000
+            minutes, nanos = divmod(nanos, nanos_per_minute)
+            if minutes:
+                if is_negative:
+                    result.append("-")
+                result.append(f"{minutes}M")
+                
+            nanos_per_second = 1000000000
+            seconds, nanos_fraction = divmod(nanos, nanos_per_second)
+            
+            if seconds or nanos_fraction:
+                if is_negative:
+                    result.append("-")
+                if seconds:
+                    result.append(str(seconds))
+                elif nanos_fraction:
+                    result.append("0")
+                    
+                if nanos_fraction:
+                    nano_str = f"{nanos_fraction:09d}"
+                    trimmed = nano_str.rstrip("0")
+                    if len(trimmed) <= 3:
+                        while len(trimmed) < 3:
+                            trimmed += "0"
+                    elif len(trimmed) <= 6:
+                        while len(trimmed) < 6:
+                            trimmed += "0"
+                    else:
+                        while len(trimmed) < 9:
+                            trimmed += "0"
+                    result.append(f".{trimmed}")
+                result.append("S")
+                
+        if len(result) == 1:
+            result.append("0Y")  # Special case for zero interval
+            
+        return "".join(result)
+
+    @classmethod
+    def from_str(cls, s: str) -> 'Interval':
+        """Parse an ISO8601 duration format string into an Interval."""
+        pattern = r'^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$'
+        match = re.match(pattern, s)
+        if not match or len(s) == 1:
+            raise ValueError(f"Invalid interval format: {s}")
+            
+        parts = match.groups()
+        if not any(parts[:3]) and not parts[3]:
+            raise ValueError(f"Invalid interval format: at least one component (Y/M/D/H/M/S) is required: {s}")
+            
+        if parts[3] == "T" and not any(parts[4:7]):
+            raise ValueError(f"Invalid interval format: time designator 'T' present but no time components specified: {s}")
+            
+        def parse_num(s: str, suffix: str) -> int:
+            if not s:
+                return 0
+            return int(s.rstrip(suffix))
+            
+        years = parse_num(parts[0], "Y")
+        months = parse_num(parts[1], "M")
+        total_months = years * 12 + months
+        
+        days = parse_num(parts[2], "D")
+        
+        nanos = 0
+        if parts[3]:  # Has time component
+            # Convert hours to nanoseconds
+            hours = parse_num(parts[4], "H")
+            nanos += hours * 3600000000000
+            
+            # Convert minutes to nanoseconds
+            minutes = parse_num(parts[5], "M")
+            nanos += minutes * 60000000000
+            
+            # Handle seconds and fractional seconds
+            if parts[6]:
+                seconds = parts[6].rstrip("S")
+                if "," in seconds:
+                    seconds = seconds.replace(",", ".")
+                    
+                if "." in seconds:
+                    sec_parts = seconds.split(".")
+                    whole_seconds = sec_parts[0] if sec_parts[0] else "0"
+                    nanos += int(whole_seconds) * 1000000000
+                    frac = sec_parts[1][:9].ljust(9, "0")
+                    frac_nanos = int(frac)
+                    if seconds.startswith("-"):
+                        frac_nanos = -frac_nanos
+                    nanos += frac_nanos
+                else:
+                    nanos += int(seconds) * 1000000000
+                    
+        return cls(months=total_months, days=days, nanos=nanos)
+
+
+@dataclass
+class NullInterval:
+    """Represents a Spanner INTERVAL that may be NULL."""
+    interval: Interval
+    valid: bool = True
+    
+    def is_null(self) -> bool:
+        return not self.valid
+        
+    def __str__(self) -> str:
+        if not self.valid:
+            return "NULL"
+        return str(self.interval)
+
+
 def _make_value_pb(value):
     """Helper for :func:`_make_list_value_pbs`.
 
@@ -251,6 +409,12 @@ def _make_value_pb(value):
             return Value(null_value="NULL_VALUE")
         else:
             return Value(string_value=base64.b64encode(value))
+    if isinstance(value, Interval):
+        return Value(string_value=str(value))
+    if isinstance(value, NullInterval):
+        if value.is_null():
+            return Value(null_value="NULL_VALUE")
+        return Value(string_value=str(value.interval))
 
     raise ValueError("Unknown type: %s" % (value,))
 
@@ -367,6 +531,8 @@ def _get_type_decoder(field_type, field_name, column_info=None):
             for item_field in field_type.struct_type.fields
         ]
         return lambda value_pb: _parse_struct(value_pb, element_decoders)
+    elif type_code == TypeCode.INTERVAL:
+        return _parse_interval
     else:
         raise ValueError("Unknown type: %s" % (field_type,))
 
@@ -473,6 +639,13 @@ def _parse_nullable(value_pb, decoder):
         return decoder(value_pb)
 
 
+def _parse_interval(value_pb):
+    """Parse a Value protobuf containing an interval."""
+    if hasattr(value_pb, 'string_value'):
+        return Interval.from_str(value_pb.string_value)
+    return Interval.from_str(value_pb)
+
+
 class _SessionWrapper(object):
     """Base class for objects wrapping a session.
 

google/cloud/spanner_v1/param_types.py

@@ -36,6 +36,7 @@
 PG_NUMERIC = Type(code=TypeCode.NUMERIC, type_annotation=TypeAnnotationCode.PG_NUMERIC)
 PG_JSONB = Type(code=TypeCode.JSON, type_annotation=TypeAnnotationCode.PG_JSONB)
 PG_OID = Type(code=TypeCode.INT64, type_annotation=TypeAnnotationCode.PG_OID)
+INTERVAL = Type(code=TypeCode.INTERVAL)
 
 
 def Array(element_type):

google/cloud/spanner_v1/streamed.py

@@ -391,6 +391,7 @@ def _merge_struct(lhs, rhs, type_):
     TypeCode.NUMERIC: _merge_string,
     TypeCode.JSON: _merge_string,
     TypeCode.PROTO: _merge_string,
+    TypeCode.INTERVAL: _merge_string,
     TypeCode.ENUM: _merge_string,
 }