Sqlalchemy compositional

.github/workflows/test.yml

@@ -72,8 +72,7 @@ jobs:
 
       - name: Initialize mypy
         run: |
-          mypy . > /dev/null || true
-          mypy --install-types --non-interactive
+          mypy --install-types --non-interactive . || true
 
       - name: Run tests
         run: |

terracotta/drivers/__init__.py

@@ -8,12 +8,14 @@
 import urllib.parse as urlparse
 from pathlib import Path
 
-from terracotta.drivers.base import Driver
+from terracotta.drivers.base import MetaStore
+from terracotta.drivers.driver import TerracottaDriver
+from terracotta.drivers.raster_base import RasterDriver
 
 URLOrPathType = Union[str, Path]
 
 
-def load_driver(provider: str) -> Type[Driver]:
+def load_driver(provider: str) -> Type[MetaStore]:
     if provider == 'sqlite-remote':
         from terracotta.drivers.sqlite_remote import RemoteSQLiteDriver
         return RemoteSQLiteDriver
@@ -42,10 +44,10 @@ def auto_detect_provider(url_or_path: Union[str, Path]) -> str:
     return 'sqlite'
 
 
-_DRIVER_CACHE: Dict[Tuple[URLOrPathType, str, int], Driver] = {}
+_DRIVER_CACHE: Dict[Tuple[URLOrPathType, str, int], TerracottaDriver] = {}
 
 
-def get_driver(url_or_path: URLOrPathType, provider: str = None) -> Driver:
+def get_driver(url_or_path: URLOrPathType, provider: str = None) -> TerracottaDriver:
     """Retrieve Terracotta driver instance for the given path.
 
     This function always returns the same instance for identical inputs.
@@ -85,6 +87,10 @@ def get_driver(url_or_path: URLOrPathType, provider: str = None) -> Driver:
     cache_key = (normalized_path, provider, os.getpid())
 
     if cache_key not in _DRIVER_CACHE:
-        _DRIVER_CACHE[cache_key] = DriverClass(url_or_path)
+        driver = TerracottaDriver(
+            metastore=DriverClass(url_or_path),
+            rasterstore=RasterDriver()
+        )
+        _DRIVER_CACHE[cache_key] = driver
 
     return _DRIVER_CACHE[cache_key]

terracotta/drivers/base.py

@@ -7,9 +7,11 @@
 import functools
 from abc import ABC, abstractmethod
 from collections import OrderedDict
-from typing import (Any, Callable, Dict, List, Mapping, Sequence, Tuple,
-                    TypeVar, Union)
+from typing import (Any, Callable, Dict, List, Mapping, Optional, Sequence,
+                    Tuple, TypeVar, Union, cast)
 
+KeysType = Mapping[str, str]
+MultiValueKeysType = Mapping[str, Union[str, List[str]]]
 Number = TypeVar('Number', int, float)
 T = TypeVar('T')
 
@@ -22,14 +24,15 @@ def requires_connection(
         return functools.partial(requires_connection, verify=verify)
 
     @functools.wraps(fun)
-    def inner(self: Driver, *args: Any, **kwargs: Any) -> T:
+    def inner(self: MetaStore, *args: Any, **kwargs: Any) -> T:
+        assert fun is not None
         with self.connect(verify=verify):
-            # Apparently mypy thinks fun might still be None, hence the ignore:
-            return fun(self, *args, **kwargs)  # type: ignore
+            return fun(self, *args, **kwargs)
+
     return inner
 
 
-class Driver(ABC):
+class MetaStore(ABC):
     """Abstract base class for all Terracotta data backends.
 
     Defines a common interface for all drivers.
@@ -105,14 +108,14 @@ def get_keys(self) -> OrderedDict:
         pass
 
     @abstractmethod
-    def get_datasets(self, where: Mapping[str, Union[str, List[str]]] = None,
+    def get_datasets(self, where: MultiValueKeysType = None,
                      page: int = 0, limit: int = None) -> Dict[Tuple[str, ...], Any]:
         # Get all known dataset key combinations matching the given constraints,
         # and a handle to retrieve the data (driver dependent)
         pass
 
     @abstractmethod
-    def get_metadata(self, keys: Union[Sequence[str], Mapping[str, str]]) -> Dict[str, Any]:
+    def get_metadata(self, keys: KeysType) -> Optional[Dict[str, Any]]:
         """Return all stored metadata for given keys.
 
         Arguments:
@@ -136,19 +139,50 @@ def get_metadata(self, keys: Union[Sequence[str], Mapping[str, str]]) -> Dict[st
         """
         pass
 
+    @abstractmethod
+    def insert(self, keys: KeysType,
+               handle: Any, **kwargs: Any) -> None:
+        """Register a new dataset. Used to populate metadata database.
+
+        Arguments:
+
+            keys: Keys of the dataset. Can either be given as a sequence of key values, or
+                as a mapping ``{key_name: key_value}``.
+            handle: Handle to access dataset (driver dependent).
+
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, keys: KeysType) -> None:
+        """Remove a dataset from the metadata database.
+
+        Arguments:
+
+            keys:  Keys of the dataset. Can either be given as a sequence of key values, or
+                as a mapping ``{key_name: key_value}``.
+
+        """
+        pass
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(\'{self.path}\')'
+
+
+class RasterStore(ABC):
+
     @abstractmethod
     # TODO: add accurate signature if mypy ever supports conditional return types
-    def get_raster_tile(self, keys: Union[Sequence[str], Mapping[str, str]], *,
+    def get_raster_tile(self, handle: str, *,
                         tile_bounds: Sequence[float] = None,
                         tile_size: Sequence[int] = (256, 256),
                         preserve_values: bool = False,
                         asynchronous: bool = False) -> Any:
-        """Load a raster tile with given keys and bounds.
+        """Load a raster tile with given handle and bounds.
 
         Arguments:
 
-            keys: Keys of the requested dataset. Can either be given as a sequence of key values,
-                or as a mapping ``{key_name: key_value}``.
+            handle: Handle of the requested dataset.
             tile_bounds: Physical bounds of the tile to read, in Web Mercator projection (EPSG3857).
                 Reads the whole dataset if not given.
             tile_size: Shape of the output array to return. Must be two-dimensional.
@@ -168,39 +202,10 @@ def get_raster_tile(self, keys: Union[Sequence[str], Mapping[str, str]], *,
         """
         pass
 
-    @staticmethod
     @abstractmethod
-    def compute_metadata(data: Any, *,
+    def compute_metadata(self, handle: str, *,
                          extra_metadata: Any = None,
-                         **kwargs: Any) -> Dict[str, Any]:
+                         use_chunks: bool = None,
+                         max_shape: Sequence[int] = None) -> Dict[str, Any]:
         # Compute metadata for a given input file (driver dependent)
         pass
-
-    @abstractmethod
-    def insert(self, keys: Union[Sequence[str], Mapping[str, str]],
-               handle: Any, **kwargs: Any) -> None:
-        """Register a new dataset. Used to populate metadata database.
-
-        Arguments:
-
-            keys: Keys of the dataset. Can either be given as a sequence of key values, or
-                as a mapping ``{key_name: key_value}``.
-            handle: Handle to access dataset (driver dependent).
-
-        """
-        pass
-
-    @abstractmethod
-    def delete(self, keys: Union[Sequence[str], Mapping[str, str]]) -> None:
-        """Remove a dataset from the metadata database.
-
-        Arguments:
-
-            keys:  Keys of the dataset. Can either be given as a sequence of key values, or
-                as a mapping ``{key_name: key_value}``.
-
-        """
-        pass
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}(\'{self.path}\')'

terracotta/drivers/driver.py

@@ -0,0 +1,177 @@
+from typing import (Any, Callable, Collection, Dict, Mapping,
+                    Sequence, Tuple, TypeVar, Union, cast)
+import contextlib
+import functools
+from collections import OrderedDict
+
+import terracotta
+from terracotta import exceptions
+from terracotta.drivers.base import (KeysType, MetaStore, MultiValueKeysType,
+                                     RasterStore, requires_connection)
+
+ExtendedKeysType = Union[Sequence[str], KeysType]
+T = TypeVar('T')
+
+
+def only_element(iterable: Collection[T]) -> T:
+    if not iterable:
+        raise exceptions.DatasetNotFoundError('No dataset found')
+    assert len(iterable) == 1
+    return next(iter(iterable))
+
+
+def standardize_keys(
+    fun: Callable[..., T] = None, *,
+    requires_all_keys: bool = False
+) -> Union[Callable[..., T], functools.partial]:
+    if fun is None:
+        return functools.partial(standardize_keys, requires_all_keys=requires_all_keys)
+
+    @functools.wraps(fun)
+    def inner(
+        self: "TerracottaDriver",
+        keys: ExtendedKeysType = None,
+        *args: Any, **kwargs: Any
+    ) -> T:
+        if requires_all_keys and (keys is None or len(keys) != len(self.key_names)):
+            raise exceptions.InvalidKeyError(
+                f'Got wrong number of keys (available keys: {self.key_names})'
+            )
+
+        if isinstance(keys, Mapping):
+            keys = dict(keys.items())
+        elif isinstance(keys, Sequence):
+            keys = dict(zip(self.key_names, keys))
+        elif keys is None:
+            keys = {}
+        else:
+            raise exceptions.InvalidKeyError(
+                'Encountered unknown key type, expected Mapping or Sequence'
+            )
+
+        unknown_keys = set(keys) - set(self.key_names)
+        if unknown_keys:
+            raise exceptions.InvalidKeyError(
+                f'Encountered unrecognized keys {unknown_keys} (available keys: {self.key_names})'
+            )
+
+        # Apparently mypy thinks fun might still be None, hence the ignore:
+        return fun(self, keys, *args, **kwargs)  # type: ignore
+    return inner
+
+
+class TerracottaDriver:
+
+    def __init__(self, metastore: MetaStore, rasterstore: RasterStore) -> None:
+        self.metastore = metastore
+        self.rasterstore = rasterstore
+
+        settings = terracotta.get_settings()
+        self.LAZY_LOADING_MAX_SHAPE: Tuple[int, int] = settings.LAZY_LOADING_MAX_SHAPE
+
+    @property
+    def db_version(self) -> str:
+        return self.metastore.db_version
+
+    @property
+    def key_names(self) -> Tuple[str, ...]:
+        return self.metastore.key_names
+
+    def create(self, keys: Sequence[str], *,
+               key_descriptions: Mapping[str, str] = None) -> None:
+        self.metastore.create(keys=keys, key_descriptions=key_descriptions)
+
+    def connect(self, verify: bool = True) -> contextlib.AbstractContextManager:
+        return self.metastore.connect(verify=verify)
+
+    @requires_connection
+    def get_keys(self) -> OrderedDict:
+        return self.metastore.get_keys()
+
+    @requires_connection
+    @standardize_keys
+    def get_datasets(self, keys: MultiValueKeysType = None,
+                     page: int = 0, limit: int = None) -> Dict[Tuple[str, ...], Any]:
+        return self.metastore.get_datasets(
+            where=keys,
+            page=page,
+            limit=limit
+        )
+
+    @requires_connection
+    @standardize_keys(requires_all_keys=True)
+    def get_metadata(self, keys: ExtendedKeysType) -> Dict[str, Any]:
+        keys = cast(KeysType, keys)
+
+        metadata = self.metastore.get_metadata(keys)
+
+        if metadata is None:
+            # metadata is not computed yet, trigger lazy loading
+            handle = only_element(self.get_datasets(keys).values())
+            metadata = self.compute_metadata(handle, max_shape=self.LAZY_LOADING_MAX_SHAPE)
+            self.insert(keys, handle, metadata=metadata)
+
+            # this is necessary to make the lazy loading tests pass...
+            metadata = self.metastore.get_metadata(keys)
+            assert metadata is not None
+
+        return metadata
+
+    @requires_connection
+    @standardize_keys(requires_all_keys=True)
+    def insert(
+        self, keys: ExtendedKeysType,
+        handle: Any, *,
+        override_path: str = None,
+        metadata: Mapping[str, Any] = None,
+        skip_metadata: bool = False,
+        **kwargs: Any
+    ) -> None:
+        keys = cast(KeysType, keys)
+
+        if metadata is None and not skip_metadata:
+            metadata = self.compute_metadata(handle)
+
+        self.metastore.insert(
+            keys=keys,
+            handle=override_path or handle,
+            metadata=metadata,
+            **kwargs
+        )
+
+    @requires_connection
+    @standardize_keys(requires_all_keys=True)
+    def delete(self, keys: ExtendedKeysType) -> None:
+        keys = cast(KeysType, keys)
+
+        self.metastore.delete(keys)
+
+    # @standardize_keys(requires_all_keys=True)
+    def get_raster_tile(self, keys: ExtendedKeysType, *,
+                        tile_bounds: Sequence[float] = None,
+                        tile_size: Sequence[int] = (256, 256),
+                        preserve_values: bool = False,
+                        asynchronous: bool = False) -> Any:
+        handle = only_element(self.get_datasets(keys).values())
+
+        return self.rasterstore.get_raster_tile(
+            handle=handle,
+            tile_bounds=tile_bounds,
+            tile_size=tile_size,
+            preserve_values=preserve_values,
+            asynchronous=asynchronous,
+        )
+
+    def compute_metadata(self, handle: str, *,
+                         extra_metadata: Any = None,
+                         use_chunks: bool = None,
+                         max_shape: Sequence[int] = None) -> Dict[str, Any]:
+        return self.rasterstore.compute_metadata(
+            handle=handle,
+            extra_metadata=extra_metadata,
+            use_chunks=use_chunks,
+            max_shape=max_shape,
+        )
+
+    def __repr__(self) -> str:
+        return self.metastore.__repr__()