Merge pull request #744 from lsst/tickets/DM-36489

DM-36489: Add pgSphere support to obscore manager

.github/workflows/build.yaml

@@ -38,10 +38,16 @@ jobs:
         run: |
           conda install -y -q sqlite
 
+      # Postgres-14 is already installed from official postgres repo, but we
+      # also need pgsphere which is not installed. The repo is not in the list,
+      # so we need to re-add it first.
       - name: Install postgresql (server)
         run: |
+          sudo sh -c \
+            'echo "deb http://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list'
+          wget --quiet -O - https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add -
           sudo apt-get update
-          sudo apt-get install postgresql
+          sudo apt-get install postgresql-14 postgresql-14-pgsphere
 
       - name: Install postgresql Python packages
         shell: bash -l {0}

python/lsst/daf/butler/core/ddl.py

@@ -38,6 +38,7 @@
     "TableSpec",
     "FieldSpec",
     "ForeignKeySpec",
+    "IndexSpec",
     "Base64Bytes",
     "Base64Region",
     "AstropyTimeNsecTai",
@@ -494,6 +495,37 @@ def fromConfig(cls, config: Config) -> ForeignKeySpec:
         )
 
 
+@dataclass(frozen=True)
+class IndexSpec:
+    """Specification of an index on table columns.
+
+    Parameters
+    ----------
+    *columns : `str`
+        Names of the columns to index.
+    **kwargs: `Any`
+        Additional keyword arguments to pass directly to
+        `sqlalchemy.schema.Index` constructor. This could be used to provide
+        backend-specific options, e.g. to create a ``GIST`` index in PostgreSQL
+        one can pass ``postgresql_using="gist"``.
+    """
+
+    def __init__(self, *columns: str, **kwargs: Any):
+        object.__setattr__(self, "columns", tuple(columns))
+        object.__setattr__(self, "kwargs", kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self.columns)
+
+    columns: Tuple[str, ...]
+    """Column names to include in the index (`Tuple` [ `str` ])."""
+
+    kwargs: dict[str, Any]
+    """Additional keyword arguments passed directly to
+    `sqlalchemy.schema.Index` constructor (`dict` [ `str`, `Any` ]).
+    """
+
+
 @dataclass
 class TableSpec:
     """A data class used to define a table or table-like query interface.
@@ -504,7 +536,7 @@ class TableSpec:
         Specifications for the columns in this table.
     unique : `Iterable` [ `tuple` [ `str` ] ], optional
         Non-primary-key unique constraints for the table.
-    indexes: `Iterable` [ `tuple` [ `str` ] ], optional
+    indexes: `Iterable` [ `IndexSpec` ], optional
         Indexes for the table.
     foreignKeys : `Iterable` [ `ForeignKeySpec` ], optional
         Foreign key constraints for the table.
@@ -527,7 +559,7 @@ def __init__(
         fields: Iterable[FieldSpec],
         *,
         unique: Iterable[Tuple[str, ...]] = (),
-        indexes: Iterable[Tuple[str, ...]] = (),
+        indexes: Iterable[IndexSpec] = (),
         foreignKeys: Iterable[ForeignKeySpec] = (),
         exclusion: Iterable[Tuple[Union[str, Type[TimespanDatabaseRepresentation]], ...]] = (),
         recycleIds: bool = True,
@@ -547,7 +579,7 @@ def __init__(
     unique: Set[Tuple[str, ...]]
     """Non-primary-key unique constraints for the table."""
 
-    indexes: Set[Tuple[str, ...]]
+    indexes: Set[IndexSpec]
     """Indexes for the table."""
 
     foreignKeys: List[ForeignKeySpec]

python/lsst/daf/butler/datastores/fileDatastore.py

@@ -237,7 +237,7 @@ def makeTableSpec(cls, datasetIdColumnType: type) -> ddl.TableSpec:
                 ddl.FieldSpec(name="file_size", dtype=BigInteger, nullable=True),
             ],
             unique=frozenset(),
-            indexes=[tuple(["path"])],
+            indexes=[ddl.IndexSpec("path")],
         )
 
     def __init__(

python/lsst/daf/butler/registry/datasets/byDimensions/tables.py

@@ -450,5 +450,5 @@ def makeCalibTableSpec(
         # in our DatasetRecordStorage.certify() implementation, and just create
         # a regular index here in the hope that helps with lookups.
         index.extend(fieldSpec.name for fieldSpec in tsFieldSpecs)
-        tableSpec.indexes.add(tuple(index))  # type: ignore
+        tableSpec.indexes.add(ddl.IndexSpec(*index))  # type: ignore
     return tableSpec

python/lsst/daf/butler/registry/dimensions/table.py

@@ -562,12 +562,12 @@ def _makeOverlapTableSpec(cls, element: DatabaseDimensionElement) -> ddl.TableSp
                 # This index has the same fields as the PK, in a different
                 # order, to facilitate queries that know skypix_index and want
                 # to find the other element.
-                (
+                ddl.IndexSpec(
                     "skypix_system",
                     "skypix_level",
                     "skypix_index",
-                )
-                + tuple(element.graph.required.names),
+                    *element.graph.required.names,
+                ),
             },
             foreignKeys=[
                 # Foreign key to summary table.  This makes sure we don't

python/lsst/daf/butler/registry/interfaces/_collections.py

@@ -621,7 +621,7 @@ def resolve_wildcard(
             their children, but not both.
 
         Returns
-        ------
+        -------
         records : `list` [ `CollectionRecord` ]
             Matching collection records.
         """

python/lsst/daf/butler/registry/interfaces/_database.py

@@ -975,14 +975,15 @@ def _convertTableSpec(
         allIndexes.update(spec.unique)
         args.extend(
             sqlalchemy.schema.Index(
-                self.shrinkDatabaseEntityName("_".join([name, "idx"] + list(columns))),
-                *columns,
-                unique=(columns in spec.unique),
+                self.shrinkDatabaseEntityName("_".join([name, "idx"] + list(index.columns))),
+                *index.columns,
+                unique=(index.columns in spec.unique),
+                **index.kwargs,
             )
-            for columns in spec.indexes
-            if columns not in allIndexes
+            for index in spec.indexes
+            if index.columns not in allIndexes
         )
-        allIndexes.update(spec.indexes)
+        allIndexes.update(index.columns for index in spec.indexes)
         args.extend(
             sqlalchemy.schema.Index(
                 self.shrinkDatabaseEntityName("_".join((name, "fkidx") + fk.source)),

python/lsst/daf/butler/registry/obscore/__init__.py

@@ -23,3 +23,4 @@
 from ._manager import *
 from ._records import *
 from ._schema import *
+from ._spatial import *

python/lsst/daf/butler/registry/obscore/_config.py

@@ -28,6 +28,7 @@
     "ExtraColumnType",
     "ObsCoreConfig",
     "ObsCoreManagerConfig",
+    "SpatialPluginConfig",
 ]
 
 import enum
@@ -60,6 +61,9 @@ class ExtraColumnConfig(BaseModel):
     length: Optional[int] = None
     """Optional length qualifier for a column, only used for strings."""
 
+    doc: Optional[str] = None
+    """Documentation string for this column."""
+
 
 class DatasetTypeConfig(BaseModel):
     """Configuration describing dataset type-related options."""
@@ -100,6 +104,16 @@ class DatasetTypeConfig(BaseModel):
     values, or ExtraColumnConfig mappings."""
 
 
+class SpatialPluginConfig(BaseModel):
+    """Configuration class for a spatial plugin."""
+
+    cls: str
+    """Name of the class implementing plugin methods."""
+
+    config: Dict[str, Any] = {}
+    """Configuration object passed to plugin ``initialize()`` method."""
+
+
 class ObsCoreConfig(BaseModel):
     """Configuration which controls conversion of Registry datasets into
     obscore records.
@@ -143,9 +157,10 @@ class ObsCoreConfig(BaseModel):
     spectral_ranges: Dict[str, Tuple[float, float]] = {}
     """Maps band name or filter name to a min/max of spectral range."""
 
-    spatial_backend: Optional[str] = None
-    """The name of a spatial backend which manages additional spatial
-    columns and indices (e.g. "pgsphere"). By default there is no spatial
+    spatial_plugins: Dict[str, SpatialPluginConfig] = {}
+    """Optional configuration for plugins managing spatial columns and
+    indices. The key is an arbitrary name and the value is an object describing
+    plugin class and its configuration options. By default there is no spatial
     indexing support, but a standard ``s_region`` column is always included.
     """