[u r] Add type and domain fields to TDRSourceSpec (#6426)

Author: nadove-ucsc

UPGRADING.rst

@@ -23,7 +23,8 @@ have too many entries in this file.
 #6426 Clean-up and generalize TDR source specs
 ==============================================
 
-The "snapshot/" string has been removed from TDR source specs.
+The "snapshot/" string has been removed from TDR source specs, and the ``type``
+and ``domain`` fields have been added.
 Update the ``mksrc`` function in ``environment.py`` for each of your personal
 deployments. As always, use the sandbox deployment's ``environment.py`` as a
 model when upgrading personal deployments.

deployments/anvilbox/environment.py

@@ -40,6 +40,8 @@ def mksrc(google_project,
         prefix = common_prefix(subgraphs)
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         prefix + '/0'

deployments/anvildev/environment.py

@@ -27,6 +27,8 @@ def mksrc(google_project, snapshot, subgraphs, flags: int = 0) -> tuple[str, str
     assert flags <= ma | pop
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         '/' + str(partition_prefix_length(subgraphs))

deployments/anvilprod/environment.py

@@ -27,6 +27,8 @@ def mksrc(google_project, snapshot, subgraphs, flags: int = 0) -> tuple[str, str
     assert flags <= ma | pop
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         '/' + str(partition_prefix_length(subgraphs))

deployments/dev/environment.py

@@ -27,6 +27,8 @@ def mksrc(google_project, snapshot, subgraphs, flags: int = 0) -> tuple[str, str
     assert flags <= ma | pop
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         '/' + str(partition_prefix_length(subgraphs))

deployments/hammerbox/environment.py

@@ -40,6 +40,8 @@ def mksrc(google_project,
         prefix = common_prefix(subgraphs)
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         prefix + '/0'

deployments/prod/environment.py

@@ -29,6 +29,8 @@ def mksrc(google_project, snapshot, subgraphs, flags: int = 0) -> tuple[str, str
     assert flags <= ma | pop
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         '/' + str(partition_prefix_length(subgraphs))

deployments/sandbox/environment.py

@@ -40,6 +40,8 @@ def mksrc(google_project,
         prefix = common_prefix(subgraphs)
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         prefix + '/0'

deployments/tempdev/environment.py

@@ -27,6 +27,8 @@ def mksrc(google_project, snapshot, subgraphs, flags: int = 0) -> tuple[str, str
     assert flags <= ma | pop
     source = None if flags & pop else ':'.join([
         'tdr',
+        'bigquery',
+        'gcp',
         google_project,
         snapshot,
         '/' + str(partition_prefix_length(subgraphs))

src/azul/terra.py

@@ -5,6 +5,7 @@
 from collections.abc import (
     Sequence,
 )
+import enum
 import json
 import logging
 from time import (
@@ -52,6 +53,7 @@
     cache,
     config,
     mutable_furl,
+    reject,
     require,
 )
 from azul.auth import (
@@ -91,37 +93,67 @@
 log = logging.getLogger(__name__)
 
 
+class SourceType(enum.Enum):
+    bigquery = 'bigquery'
+    parquet = 'parquet'
+
+
+class SourceDomain(enum.Enum):
+    gcp = 'gcp'
+    azure = 'azure'
+
+
 @attrs.frozen(kw_only=True)
 class TDRSourceSpec(SourceSpec):
+    type: SourceType
+    domain: SourceDomain
     subdomain: str
     name: str
 
     @classmethod
     def parse(cls, spec: str) -> 'TDRSourceSpec':
         """
         Construct an instance from its string representation, using the syntax
-        'tdr:{subdomain}:{name}:{prefix}' ending with an optional
+        'tdr:{type}{domain}{subdomain}:{name}:{prefix}' ending with an optional
         '/{partition_prefix_length}'.
 
-        >>> s = TDRSourceSpec.parse('tdr:foo:bar:/0')
+        >>> s = TDRSourceSpec.parse('tdr:bigquery:gcp:foo:bar:/0')
         >>> s # doctest: +NORMALIZE_WHITESPACE
         TDRSourceSpec(prefix=Prefix(common='', partition=0),
+                      type=<SourceType.bigquery: 'bigquery'>,
+                      domain=<SourceDomain.gcp: 'gcp'>,
                       subdomain='foo',
                       name='bar')
 
         >>> str(s)
-        'tdr:foo:bar:/0'
+        'tdr:bigquery:gcp:foo:bar:/0'
+
+        >>> TDRSourceSpec.parse('tdr:spam:gcp:foo:bar:/0')
+        Traceback (most recent call last):
+        ...
+        ValueError: 'spam' is not a valid SourceType
+
+        >>> TDRSourceSpec.parse('tdr:bigquery:eggs:foo:bar:/0')
+        Traceback (most recent call last):
+        ...
+        ValueError: 'eggs' is not a valid SourceDomain
 
-        >>> TDRSourceSpec.parse('tdr:foo:bar:n32/0')
+        >>> TDRSourceSpec.parse('tdr:bigquery:gcp:foo:bar:n32/0')
         Traceback (most recent call last):
         ...
         azul.uuids.InvalidUUIDPrefixError: 'n32' is not a valid UUID prefix.
         """
         rest, prefix = cls._parse(spec)
         # BigQuery (and by extension the TDR) does not allow : or / in dataset names
-        service, subdomain, name = rest.split(':')
+        service, type, domain, subdomain, name = rest.split(':')
         assert service == 'tdr', service
+        type = SourceType(type)
+        reject(type == SourceType.parquet, 'Parquet sources are not yet supported')
+        domain = SourceDomain(domain)
+        reject(domain == SourceDomain.azure, 'Azure sources are not yet supported')
         self = cls(prefix=prefix,
+                   type=SourceType(type),
+                   domain=SourceDomain(domain),
                    subdomain=subdomain,
                    name=name)
         assert spec == str(self), spec
@@ -131,20 +163,22 @@ def __str__(self) -> str:
         """
         The inverse of :meth:`parse`.
 
-        >>> s = 'tdr:foo:bar:/0'
+        >>> s = 'tdr:bigquery:gcp:foo:bar:/0'
         >>> s == str(TDRSourceSpec.parse(s))
         True
 
-        >>> s = 'tdr:foo:bar:22/0'
+        >>> s = 'tdr:bigquery:gcp:foo:bar:22/0'
         >>> s == str(TDRSourceSpec.parse(s))
         True
 
-        >>> s = 'tdr:foo:bar:22/2'
+        >>> s = 'tdr:bigquery:gcp:foo:bar:22/2'
         >>> s == str(TDRSourceSpec.parse(s))
         True
         """
         return ':'.join([
             'tdr',
+            self.type.value,
+            self.domain.value,
             self.subdomain,
             self.name,
             str(self.prefix)
@@ -157,18 +191,20 @@ def contains(self, other: 'SourceSpec') -> bool:
         """
         >>> p = TDRSourceSpec.parse
 
-        >>> p('tdr:foo:bar:/0').contains(p('tdr:foo:bar:/0'))
+        >>> p('tdr:bigquery:gcp:foo:bar:/0').contains(p('tdr:bigquery:gcp:foo:bar:/0'))
         True
 
-        >>> p('tdr:foo:bar:/0').contains(p('tdr:bar:bar:/0'))
+        >>> p('tdr:bigquery:gcp:foo:bar:/0').contains(p('tdr:bigquery:gcp:bar:bar:/0'))
         False
 
-        >>> p('tdr:foo:bar:/0').contains(p('tdr:foo:baz:/0'))
+        >>> p('tdr:bigquery:gcp:foo:bar:/0').contains(p('tdr:bigquery:gcp:foo:baz:/0'))
         False
         """
         return (
             isinstance(other, TDRSourceSpec)
             and super().contains(other)
+            and self.type == other.type
+            and self.domain == other.domain
             and self.subdomain == other.subdomain
             and self.name == other.name
         )