Merge pull request #915 from lsst/tickets/DM-41158

DM-41158: (mostly) implement QueryDriver for DirectButler

python/lsst/daf/butler/__init__.py

@@ -75,12 +75,9 @@
 
 # Do not import or lift symbols from mapping_factory and persistence_content,
 # as those are internal.
-# Only lift 'Progress' from 'progess'; the module is imported as-is above
+# Only lift 'Progress' from 'progress'; the module is imported as-is above
 from .progress import Progress
 
-# Only import the main public symbols from queries
-from .queries import DataCoordinateQueryResults, DatasetRefQueryResults, DimensionRecordQueryResults, Query
-
 # Do not import or lift symbols from 'server' or 'server_models'.
 # Import the registry subpackage directly for other symbols.
 from .registry import (
@@ -93,6 +90,9 @@
 from .transfers import RepoExportContext, YamlRepoExportBackend, YamlRepoImportBackend
 from .version import *
 
+# Do not import or lift symbols from 'queries' until they are public.
+
+
 # Do not import or lift symbols from 'repo_relocation'.
 
 # Do not export the utility routines from utils.

python/lsst/daf/butler/column_spec.py

@@ -53,6 +53,7 @@
 
 from . import arrow_utils, ddl
 from ._timespan import Timespan
+from .name_shrinker import NameShrinker
 
 ColumnType: TypeAlias = Literal[
     "int", "string", "hash", "float", "datetime", "bool", "uuid", "timespan", "region"
@@ -79,11 +80,14 @@ class _BaseColumnSpec(pydantic.BaseModel, ABC):
         description="Whether the column may be ``NULL``.",
     )
 
-    def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:
+    def to_sql_spec(self, name_shrinker: NameShrinker | None = None, **kwargs: Any) -> ddl.FieldSpec:
         """Convert this specification to a SQL-specific one.
 
         Parameters
         ----------
+        name_shrinker : `NameShrinker`, optional
+            Object that should be used to shrink the field name to ensure it
+            fits within database-specific limits.
         **kwargs
             Forwarded to `ddl.FieldSpec`.
 
@@ -92,7 +96,10 @@ def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:
         sql_spec : `ddl.FieldSpec`
             A SQL-specific version of this specification.
         """
-        return ddl.FieldSpec(name=self.name, dtype=ddl.VALID_CONFIG_COLUMN_TYPES[self.type], **kwargs)
+        name = self.name
+        if name_shrinker is not None:
+            name = name_shrinker.shrink(name)
+        return ddl.FieldSpec(name=name, dtype=ddl.VALID_CONFIG_COLUMN_TYPES[self.type], **kwargs)
 
     @abstractmethod
     def to_arrow(self) -> arrow_utils.ToArrow:
@@ -162,9 +169,9 @@ class StringColumnSpec(_BaseColumnSpec):
     length: int
     """Maximum length of strings."""
 
-    def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:
+    def to_sql_spec(self, name_shrinker: NameShrinker | None = None, **kwargs: Any) -> ddl.FieldSpec:
         # Docstring inherited.
-        return super().to_sql_spec(length=self.length, **kwargs)
+        return super().to_sql_spec(length=self.length, name_shrinker=name_shrinker, **kwargs)
 
     def to_arrow(self) -> arrow_utils.ToArrow:
         # Docstring inherited.
@@ -182,9 +189,9 @@ class HashColumnSpec(_BaseColumnSpec):
     nbytes: int
     """Number of bytes for the hash."""
 
-    def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:
+    def to_sql_spec(self, name_shrinker: NameShrinker | None = None, **kwargs: Any) -> ddl.FieldSpec:
         # Docstring inherited.
-        return super().to_sql_spec(nbytes=self.nbytes, **kwargs)
+        return super().to_sql_spec(nbytes=self.nbytes, name_shrinker=name_shrinker, **kwargs)
 
     def to_arrow(self) -> arrow_utils.ToArrow:
         # Docstring inherited.

python/lsst/daf/butler/direct_butler.py

@@ -66,6 +66,7 @@
 from ._timespan import Timespan
 from .datastore import Datastore, NullDatastore
 from .dimensions import DataCoordinate, Dimension
+from .direct_query_driver import DirectQueryDriver
 from .progress import Progress
 from .queries import Query
 from .registry import (
@@ -2136,7 +2137,13 @@ def dimensions(self) -> DimensionUniverse:
     @contextlib.contextmanager
     def _query(self) -> Iterator[Query]:
         # Docstring inherited.
-        raise NotImplementedError("TODO DM-41159")
+        driver = DirectQueryDriver(
+            self._registry._db, self.dimensions, self._registry._managers, self._registry.defaults
+        )
+        query = Query(driver)
+        with self._caching_context():
+            with driver:
+                yield query
 
     def _preload_cache(self) -> None:
         """Immediately load caches that are used for common operations."""

python/lsst/daf/butler/direct_query_driver/__init__.py

@@ -0,0 +1,30 @@
+# This file is part of daf_butler.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This software is dual licensed under the GNU General Public License and also
+# under a 3-clause BSD license. Recipients may choose which of these licenses
+# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
+# respectively.  If you choose the GPL option then the following text applies
+# (but note that there is still no warranty even if you opt for BSD instead):
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from ._driver import DirectQueryDriver
+from ._postprocessing import Postprocessing
+from ._query_builder import QueryBuilder, QueryJoiner