Skip to content

Commit

Permalink
Merge pull request #27 from lsst-dm/tickets/DM-41570
Browse files Browse the repository at this point in the history
DM-41570: Update imports for new location of SqlRegistry class
  • Loading branch information
andy-slac authored Nov 3, 2023
2 parents 6f53698 + baadca2 commit 4f4f978
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 26 deletions.
8 changes: 4 additions & 4 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -1,26 +1,26 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.4.0
rev: v4.5.0
hooks:
- id: check-yaml
- id: end-of-file-fixer
- id: trailing-whitespace
- repo: https://github.com/psf/black
rev: 23.3.0
rev: 23.10.1
hooks:
- id: black
# It is recommended to specify the latest version of Python
# supported by your project here, or alternatively use
# pre-commit's default_language_version, see
# https://pre-commit.com/#top_level-default_language_version
language_version: python3.10
language_version: python3.11
- repo: https://github.com/pycqa/isort
rev: 5.12.0
hooks:
- id: isort
name: isort (python)
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.0.277
rev: v0.1.3
hooks:
- id: ruff
2 changes: 1 addition & 1 deletion migrations/obscore-config/4fe28ef5030f.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

import yaml
from alembic import context, op
from lsst.daf.butler.registries.sql import SqlRegistry
from lsst.daf.butler.registry.sql_registry import SqlRegistry
from lsst.daf.butler_migrate.butler_attributes import ButlerAttributes
from lsst.daf.butler_migrate.registry import make_registry
from lsst.utils import doImportType
Expand Down
6 changes: 3 additions & 3 deletions python/lsst/daf/butler_migrate/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,10 @@

from __future__ import annotations

from lsst.daf.butler import Registry
from lsst.daf.butler.registry.sql_registry import SqlRegistry


def make_registry(repository: str, writeable: bool = True) -> Registry:
def make_registry(repository: str, writeable: bool = True) -> SqlRegistry:
"""Make Registry instance.
Parameters
Expand All @@ -34,4 +34,4 @@ def make_registry(repository: str, writeable: bool = True) -> Registry:
writeable : `bool`, optional
If `True` (default) create a read-write connection to the database.
"""
return Registry.fromConfig(config=repository, writeable=writeable)
return SqlRegistry.fromConfig(config=repository, writeable=writeable)
39 changes: 21 additions & 18 deletions python/lsst/daf/butler_migrate/script/rewrite_sqlite_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,10 @@

from lsst.daf.butler import Butler, Config, DatasetAssociation, DatasetId, DatasetRef, SkyPixDimension
from lsst.daf.butler.datastores.fileDatastore import FileDatastore
from lsst.daf.butler.registries.sql import SqlRegistry
from lsst.daf.butler.direct_butler import DirectButler
from lsst.daf.butler.registry import CollectionType
from lsst.daf.butler.registry.databases.sqlite import SqliteDatabase
from lsst.daf.butler.registry.sql_registry import SqlRegistry
from lsst.daf.butler.transfers import RepoExportContext
from lsst.resources import ResourcePath
from lsst.utils.introspection import get_class_of
Expand All @@ -56,11 +57,12 @@ def rewrite_sqlite_registry(source: str) -> None:
"""
# Create the source butler early so we can ask it questions
# without assuming things.
source_butler = Butler(source, writeable=False)
assert isinstance(source_butler.registry, SqlRegistry), "Expecting SqlRegistry instance"
source_butler = Butler.from_config(source, writeable=False)
assert isinstance(source_butler, DirectButler)
assert isinstance(source_butler._registry, SqlRegistry), "Expecting SqlRegistry instance"

# Check that we are really working with a SQLite database.
if not isinstance(source_butler.registry._db, SqliteDatabase):
if not isinstance(source_butler._registry._db, SqliteDatabase):
raise RuntimeError("This command can only be used on SQLite registries.")

# The source butler knows where its config came from.
Expand Down Expand Up @@ -120,15 +122,16 @@ def rewrite_sqlite_registry(source: str) -> None:
)

# Create destination butler
dest_butler = Butler(dest_config, writeable=True)
assert isinstance(dest_butler.registry, SqlRegistry), "Expecting SqlRegistry instance"
assert isinstance(dest_butler.registry._db, SqliteDatabase), "Expecting SqliteDatabase instance"
dest_butler = Butler.from_config(dest_config, writeable=True)
assert isinstance(dest_butler, DirectButler)
assert isinstance(dest_butler._registry, SqlRegistry), "Expecting SqlRegistry instance"
assert isinstance(dest_butler._registry._db, SqliteDatabase), "Expecting SqliteDatabase instance"

transfer_everything(source_butler, dest_butler)

# Obtain the name of the sqlite file at the destination.
assert dest_butler.registry._db.filename is not None, "Expecting non-None filename from registry"
dest_registry_uri = ResourcePath(dest_butler.registry._db.filename)
assert dest_butler._registry._db.filename is not None, "Expecting non-None filename from registry"
dest_registry_uri = ResourcePath(dest_butler._registry._db.filename)

# Finished with writing to the destination butler so
# delete the variable to ensure we can't do any more.
Expand All @@ -138,8 +141,8 @@ def rewrite_sqlite_registry(source: str) -> None:
# and move the existing registry to a backup.

# Relocate the source registry first
assert source_butler.registry._db.filename is not None, "Expecting non-None filename from registry"
source_registry_uri = ResourcePath(source_butler.registry._db.filename)
assert source_butler._registry._db.filename is not None, "Expecting non-None filename from registry"
source_registry_uri = ResourcePath(source_butler._registry._db.filename)
new_basename = "original_" + source_registry_uri.basename()
backup_registry_uri = source_registry_uri.updatedFile(new_basename)
os.rename(source_registry_uri.ospath, backup_registry_uri.ospath)
Expand All @@ -160,16 +163,16 @@ def rewrite_sqlite_registry(source: str) -> None:
print(f"Successfully rewrote registry for butler at {source_config_uri}")


def transfer_everything(source_butler: Butler, dest_butler: Butler) -> None:
def transfer_everything(source_butler: DirectButler, dest_butler: DirectButler) -> None:
"""Transfer all content from one butler to another butler.
Assumes that both registries have a common dimension universe.
Parameters
----------
source_butler : `~lsst.daf.butler.Butler`
source_butler : `~lsst.daf.butler.direct_butler.DirectButler`
Butler to use as a source of information.
dest_butler : `~lsst.daf.butler.Butler`
dest_butler : `~lsst.daf.butler.direct_butler.DirectButler`
Butler to receive all the content.
"""
# Read all the datasets we are going to transfer, removing duplicates.
Expand Down Expand Up @@ -237,14 +240,14 @@ def create_associations(
)


def transfer_non_datasets(source_butler: Butler, dest_butler: Butler) -> None:
def transfer_non_datasets(source_butler: DirectButler, dest_butler: DirectButler) -> None:
"""Transfer everything that isn't related to datasets.
Parameters
----------
source_butler : `~lsst.daf.butler.Butler`
source_butler : `~lsst.daf.butler.direct_butler.DirectButler`
Butler to extract information from.
dest_butler : `~lsst.daf.butler.Butler`
dest_butler : `~lsst.daf.butler.direct_butler.DirectButler`
Destination butler.
"""
# Use a string buffer to save on file I/O. This might result in twice
Expand All @@ -256,7 +259,7 @@ def transfer_non_datasets(source_butler: Butler, dest_butler: Butler) -> None:
BackendClass = get_class_of(source_butler._config["repo_transfer_formats", "yaml", "export"])
backend = BackendClass(yamlBuffer)
exporter = RepoExportContext(
source_butler.registry, source_butler.datastore, backend, directory=None, transfer=None
source_butler._registry, source_butler.datastore, backend, directory=None, transfer=None
)

# Export all the collections.
Expand Down

0 comments on commit 4f4f978

Please sign in to comment.