Merge lp:~mwhudson/launchpad/code-import-optimizations into lp:launchpad

[Michael Hudson-Doyle]

Improve logging during code imports and don't build a working tree if not needed, both of which will help with large imports timing out on startup

Hi,

This branch should fix some very large foreign branch imports timing out in two ways:

1) There was no progress reporting during fetching the branch from the store or building the tree, so I moved the initial fetch into the block that has our custom ui factory installed.

2) Building the working tree was a waste of time anyway for foreign branch imports, so we no longer do that. This required the bulk of the changes from a textual point of view, to talk about branches where previously we talked about trees, but I think the result is better even if it wasn't faster.

I also add a few logging calls as suggested in bug 559808.

Cheers,
mwh

[Tim Penhey (thumper)]

The change looks fine.

It does get me wondering a bit why we didn't seem to get any logging when it was pushing the branch to the central store before. Thanks for adding the job started and finished logging bits too.

I agree that it makes more sense to use branch rather than working tree in the workers.

review: Approve

[Michael Hudson-Doyle (mwhudson)]

On 12/04/10 14:16, Tim Penhey wrote:
> Review: Approve
> The change looks fine.

Thanks.

> It does get me wondering a bit why we didn't seem to get any logging when it was pushing the branch to the central store before. Thanks for adding the job started and finished logging bits too.

I think we do actually get logging when pushing to the central store,
but the new log lines should let us see what comes from where :-)

> I agree that it makes more sense to use branch rather than working tree in the workers.

Cool.

Cheers,
mwh

=== modified file 'lib/lp/codehosting/codeimport/tests/test_worker.py'
--- lib/lp/codehosting/codeimport/tests/test_worker.py	2010-04-09 04:39:38 +0000
+++ lib/lp/codehosting/codeimport/tests/test_worker.py	2010-04-12 01:54:17 +0000
@@ -19,7 +19,7 @@
 from bzrlib.tests import TestCaseWithTransport
 from bzrlib.transport import get_transport
 from bzrlib.upgrade import upgrade
-from bzrlib.urlutils import join as urljoin
+from bzrlib.urlutils import join as urljoin, local_path_from_url
 
 from CVS import Repository, tree as CVSTree
 
@@ -120,22 +120,57 @@
 
     def test_getNewBranch(self):
         # If there's no Bazaar branch of this id, then pull creates a new
-        # Bazaar working tree.
+        # Bazaar branch.
         store = self.makeBranchStore()
-        bzr_working_tree = store.pull(
+        bzr_branch = store.pull(
             self.arbitrary_branch_id, self.temp_dir, default_format)
-        self.assertEqual([], bzr_working_tree.branch.revision_history())
+        self.assertEqual([], bzr_branch.revision_history())
+
+    def test_getNewBranch_without_tree(self):
+        # If pull() with needs_tree=False creates a new branch, it doesn't
+        # create a working tree.
+        store = self.makeBranchStore()
+        bzr_branch = store.pull(
+            self.arbitrary_branch_id, self.temp_dir, default_format, False)
+        self.assertFalse(bzr_branch.bzrdir.has_workingtree())
+
+    def test_getNewBranch_with_tree(self):
+        # If pull() with needs_tree=True creates a new branch, it creates a
+        # working tree.
+        store = self.makeBranchStore()
+        bzr_branch = store.pull(
+            self.arbitrary_branch_id, self.temp_dir, default_format, True)
+        self.assertTrue(bzr_branch.bzrdir.has_workingtree())
 
     def test_pushBranchThenPull(self):
         # After we've pushed up a branch to the store, we can then pull it
         # from the store.
         store = self.makeBranchStore()
         tree = create_branch_with_one_revision('original')
-        store.push(self.arbitrary_branch_id, tree, default_format)
-        new_tree = store.pull(
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        new_branch = store.pull(
             self.arbitrary_branch_id, self.temp_dir, default_format)
         self.assertEqual(
-            tree.branch.last_revision(), new_tree.branch.last_revision())
+            tree.branch.last_revision(), new_branch.last_revision())
+
+    def test_pull_without_needs_tree_doesnt_create_tree(self):
+        # pull with needs_tree=False doesn't spend the time to create a
+        # working tree.
+        store = self.makeBranchStore()
+        tree = create_branch_with_one_revision('original')
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        new_branch = store.pull(
+            self.arbitrary_branch_id, self.temp_dir, default_format, False)
+        self.assertFalse(new_branch.bzrdir.has_workingtree())
+
+    def test_pull_needs_tree_creates_tree(self):
+        # pull with needs_tree=True creates a working tree.
+        store = self.makeBranchStore()
+        tree = create_branch_with_one_revision('original')
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        new_branch = store.pull(
+            self.arbitrary_branch_id, self.temp_dir, default_format, True)
+        self.assertTrue(new_branch.bzrdir.has_workingtree())
 
     # XXX Tim Penhey 2009-09-18 bug 432217 Automatic upgrade of import
     # branches disabled.  Need an orderly upgrade process.
@@ -188,21 +223,21 @@
         # store.
         store = self.makeBranchStore()
         tree = create_branch_with_one_revision('original')
-        store.push(self.arbitrary_branch_id, tree, default_format)
-        store.push(self.arbitrary_branch_id, tree, default_format)
-        new_tree = store.pull(
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        new_branch = store.pull(
             self.arbitrary_branch_id, self.temp_dir, default_format)
         self.assertEqual(
-            tree.branch.last_revision(), new_tree.branch.last_revision())
+            tree.branch.last_revision(), new_branch.last_revision())
 
     def test_push_divergant_branches(self):
         # push() uses overwrite=True, so divergent branches (rebased) can be
         # pushed.
         store = self.makeBranchStore()
         tree = create_branch_with_one_revision('original')
-        store.push(self.arbitrary_branch_id, tree, default_format)
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
         tree = create_branch_with_one_revision('divergant')
-        store.push(self.arbitrary_branch_id, tree, default_format)
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
 
     def fetchBranch(self, from_url, target_path):
         """Pull a branch from `from_url` to `target_path`.
@@ -221,7 +256,7 @@
         # doesn't already exist.
         store = BazaarBranchStore(self.get_transport('doesntexist'))
         tree = create_branch_with_one_revision('original')
-        store.push(self.arbitrary_branch_id, tree, default_format)
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
         self.assertIsDirectory('doesntexist', self.get_transport())
 
     def test_storedLocation(self):
@@ -229,12 +264,12 @@
         # the BazaarBranchStore's transport.
         store = self.makeBranchStore()
         tree = create_branch_with_one_revision('original')
-        store.push(self.arbitrary_branch_id, tree, default_format)
-        new_tree = self.fetchBranch(
+        store.push(self.arbitrary_branch_id, tree.branch, default_format)
+        new_branch = self.fetchBranch(
             urljoin(store.transport.base, '%08x' % self.arbitrary_branch_id),
             'new_tree')
         self.assertEqual(
-            tree.branch.last_revision(), new_tree.branch.last_revision())
+            tree.branch.last_revision(), new_branch.last_revision())
 
     def test_sftpPrefix(self):
         # Since branches are mirrored by importd via sftp, _getMirrorURL must
@@ -270,15 +305,14 @@
             None, None, [('add', ('', 'root-id', 'directory', ''))])
         revid1 = builder.build_snapshot(None, [revid], [])
         revid2 = builder.build_snapshot(None, [revid], [])
-        branch = builder.get_branch()
-        source_tree = branch.bzrdir.create_workingtree()
         store = self.makeBranchStore()
-        store.push(self.arbitrary_branch_id, source_tree, default_format)
-        retrieved_tree = store.pull(
+        store.push(
+            self.arbitrary_branch_id, builder.get_branch(), default_format)
+        retrieved_branch = store.pull(
             self.arbitrary_branch_id, 'pulled', default_format)
         self.assertEqual(
             set([revid, revid1, revid2]),
-            set(retrieved_tree.branch.repository.all_revision_ids()))
+            set(retrieved_branch.repository.all_revision_ids()))
 
 
 class TestImportDataStore(WorkerTest):
@@ -532,21 +566,21 @@
         worker = self.makeImportWorker()
         self.assertEqual(self.source_details, worker.source_details)
 
-    def test_getBazaarWorkingTreeMakesEmptyTree(self):
-        # getBazaarWorkingTree returns a brand-new working tree for an initial
+    def test_getBazaarWorkingBranchMakesEmptyBranch(self):
+        # getBazaarBranch returns a brand-new working tree for an initial
         # import.
         worker = self.makeImportWorker()
-        bzr_working_tree = worker.getBazaarWorkingTree()
-        self.assertEqual([], bzr_working_tree.branch.revision_history())
+        bzr_branch = worker.getBazaarBranch()
+        self.assertEqual([], bzr_branch.revision_history())
 
-    def test_bazaarWorkingTreeLocation(self):
-        # getBazaarWorkingTree makes the working tree under the current
-        # working directory.
+    def test_bazaarBranchLocation(self):
+        # getBazaarBranch makes the working tree under the current working
+        # directory.
         worker = self.makeImportWorker()
-        bzr_working_tree = worker.getBazaarWorkingTree()
+        bzr_branch = worker.getBazaarBranch()
         self.assertIsSameRealPath(
-            os.path.abspath(worker.BZR_WORKING_TREE_PATH),
-            os.path.abspath(bzr_working_tree.basedir))
+            os.path.abspath(worker.BZR_BRANCH_PATH),
+            os.path.abspath(local_path_from_url(bzr_branch.base)))
 
 
 class TestCSCVSWorker(WorkerTest):
@@ -591,22 +625,22 @@
             source_details, self.get_transport('import_data'),
             self.makeBazaarBranchStore(), logging.getLogger("silent"))
 
-    def test_pushBazaarWorkingTree_saves_git_cache(self):
-        # GitImportWorker.pushBazaarWorkingTree saves a tarball of the git
-        # cache from the tree's repository in the worker's ImportDataStore.
+    def test_pushBazaarBranch_saves_git_cache(self):
+        # GitImportWorker.pushBazaarBranch saves a tarball of the git cache
+        # from the tree's repository in the worker's ImportDataStore.
         content = self.factory.getUniqueString()
-        tree = self.make_branch_and_tree('.')
-        tree.branch.repository._transport.mkdir('git')
-        tree.branch.repository._transport.put_bytes('git/cache', content)
+        branch = self.make_branch('.')
+        branch.repository._transport.mkdir('git')
+        branch.repository._transport.put_bytes('git/cache', content)
         import_worker = self.makeImportWorker()
-        import_worker.pushBazaarWorkingTree(tree)
+        import_worker.pushBazaarBranch(branch)
         import_worker.import_data_store.fetch('git-cache.tar.gz')
         extract_tarball('git-cache.tar.gz', '.')
         self.assertEqual(content, open('cache').read())
 
-    def test_getBazaarWorkingTree_fetches_legacy_git_db(self):
-        # GitImportWorker.getBazaarWorkingTree fetches the legacy git.db file,
-        # if present, from the worker's ImportDataStore into the tree's
+    def test_getBazaarBranch_fetches_legacy_git_db(self):
+        # GitImportWorker.getBazaarBranch fetches the legacy git.db file, if
+        # present, from the worker's ImportDataStore into the tree's
         # repository.
         import_worker = self.makeImportWorker()
         # Store the git.db file in the store.
@@ -614,15 +648,15 @@
         open('git.db', 'w').write(content)
         import_worker.import_data_store.put('git.db')
         # Make sure there's a Bazaar branch in the branch store.
-        tree = self.make_branch_and_tree('tree')
-        ImportWorker.pushBazaarWorkingTree(import_worker, tree)
+        branch = self.make_branch('branch')
+        ImportWorker.pushBazaarBranch(import_worker, branch)
         # Finally, fetching the tree gets the git.db file too.
-        tree = import_worker.getBazaarWorkingTree()
+        branch = import_worker.getBazaarBranch()
         self.assertEqual(
-            content, tree.branch.repository._transport.get('git.db').read())
+            content, branch.repository._transport.get('git.db').read())
 
-    def test_getBazaarWorkingTree_fetches_git_cache(self):
-        # GitImportWorker.getBazaarWorkingTree fetches the tarball of the git
+    def test_getBazaarBranch_fetches_git_cache(self):
+        # GitImportWorker.getBazaarBranch fetches the tarball of the git
         # cache from the worker's ImportDataStore and expands it into the
         # tree's repository.
         import_worker = self.makeImportWorker()
@@ -633,12 +667,13 @@
         create_tarball('cache', 'git-cache.tar.gz')
         import_worker.import_data_store.put('git-cache.tar.gz')
         # Make sure there's a Bazaar branch in the branch store.
-        tree = self.make_branch_and_tree('tree')
-        ImportWorker.pushBazaarWorkingTree(import_worker, tree)
+        branch = self.make_branch('branch')
+        ImportWorker.pushBazaarBranch(import_worker, branch)
         # Finally, fetching the tree gets the git.db file too.
-        tree = import_worker.getBazaarWorkingTree()
+        new_branch = import_worker.getBazaarBranch()
         self.assertEqual(
-            content, tree.branch.repository._transport.get('git/git-cache').read())
+            content,
+            new_branch.repository._transport.get('git/git-cache').read())
 
 
 def clean_up_default_stores_for_import(source_details):
 
=== modified file 'lib/lp/codehosting/codeimport/worker.py'
--- lib/lp/codehosting/codeimport/worker.py	2010-04-09 04:39:38 +0000
+++ lib/lp/codehosting/codeimport/worker.py	2010-04-12 01:54:17 +0000
@@ -64,17 +64,22 @@
         """Return the URL that `db_branch` is stored at."""
         return urljoin(self.transport.base, '%08x' % db_branch_id)
 
-    def pull(self, db_branch_id, target_path, required_format):
-        """Pull down the Bazaar branch for `code_import` to `target_path`.
+    def pull(self, db_branch_id, target_path, required_format,
+             needs_tree=False):
+        """Pull down the Bazaar branch of an import to `target_path`.
 
-        :return: A Bazaar working tree for the branch of `code_import`.
+        :return: A Bazaar branch for the code import corresponding to the
+            database branch with id `db_branch_id`.
         """
         remote_url = self._getMirrorURL(db_branch_id)
         try:
             remote_bzr_dir = BzrDir.open(remote_url)
         except NotBranchError:
-            return BzrDir.create_standalone_workingtree(
-                target_path, required_format)
+            local_branch = BzrDir.create_branch_and_repo(
+                target_path, format=required_format)
+            if needs_tree:
+                local_branch.bzrdir.create_workingtree()
+            return local_branch
         # XXX Tim Penhey 2009-09-18 bug 432217 Automatic upgrade of import
         # branches disabled.  Need an orderly upgrade process.
         if False and remote_bzr_dir.needs_format_conversion(
@@ -84,33 +89,33 @@
             except NoSuchFile:
                 pass
             upgrade(remote_url, required_format)
-        local_bzr_dir = remote_bzr_dir.sprout(target_path)
+        local_branch = remote_bzr_dir.sprout(
+            target_path, create_tree_if_local=needs_tree).open_branch()
         # Because of the way we do incremental imports, there may be revisions
         # in the branch's repo that are not in the ancestry of the branch tip.
         # We need to transfer them too.
-        local_bzr_dir.open_repository().fetch(
+        local_branch.repository.fetch(
             remote_bzr_dir.open_repository())
-        return local_bzr_dir.open_workingtree()
+        return local_branch
 
-    def push(self, db_branch_id, bzr_tree, required_format):
-        """Push up `bzr_tree` as the Bazaar branch for `code_import`.
+    def push(self, db_branch_id, bzr_branch, required_format):
+        """Push up `bzr_branch` as the Bazaar branch for `code_import`.
 
         :return: A boolean that is true if the push was non-trivial
             (i.e. actually transferred revisions).
         """
         self.transport.create_prefix()
-        branch_from = bzr_tree.branch
         target_url = self._getMirrorURL(db_branch_id)
         try:
-            branch_to = Branch.open(target_url)
+            remote_branch = Branch.open(target_url)
         except NotBranchError:
-            branch_to = BzrDir.create_branch_and_repo(
+            remote_branch = BzrDir.create_branch_and_repo(
                 target_url, format=required_format)
-        pull_result = branch_to.pull(branch_from, overwrite=True)
+        pull_result = remote_branch.pull(bzr_branch, overwrite=True)
         # Because of the way we do incremental imports, there may be revisions
         # in the branch's repo that are not in the ancestry of the branch tip.
         # We need to transfer them too.
-        branch_to.repository.fetch(branch_from.repository)
+        remote_branch.repository.fetch(bzr_branch.repository)
         return pull_result.old_revid != pull_result.new_revid
 
 
@@ -352,7 +357,10 @@
     """Oversees the actual work of a code import."""
 
     # Where the Bazaar working tree will be stored.
-    BZR_WORKING_TREE_PATH = 'bzr_working_tree'
+    BZR_BRANCH_PATH = 'bzr_branch'
+
+    # Should `getBazaarBranch` create a working tree?
+    needs_bzr_tree = True
 
     required_format = BzrDirFormat.get_default_format()
 
@@ -372,21 +380,22 @@
             import_data_transport, self.source_details)
         self._logger = logger
 
-    def getBazaarWorkingTree(self):
-        """Return the Bazaar `WorkingTree` that we are importing into."""
-        if os.path.isdir(self.BZR_WORKING_TREE_PATH):
-            shutil.rmtree(self.BZR_WORKING_TREE_PATH)
+    def getBazaarBranch(self):
+        """Return the Bazaar `Branch` that we are importing into."""
+        if os.path.isdir(self.BZR_BRANCH_PATH):
+            shutil.rmtree(self.BZR_BRANCH_PATH)
         return self.bazaar_branch_store.pull(
-            self.source_details.branch_id, self.BZR_WORKING_TREE_PATH,
-            self.required_format)
+            self.source_details.branch_id, self.BZR_BRANCH_PATH,
+            self.required_format, self.needs_bzr_tree)
 
-    def pushBazaarWorkingTree(self, bazaar_tree):
-        """Push the updated Bazaar working tree to the server.
+    def pushBazaarBranch(self, bazaar_branch):
+        """Push the updated Bazaar branch to the server.
 
         :return: True if revisions were transferred.
         """
         return self.bazaar_branch_store.push(
-            self.source_details.branch_id, bazaar_tree, self.required_format)
+            self.source_details.branch_id, bazaar_branch,
+            self.required_format)
 
     def getWorkingDirectory(self):
         """The directory we should change to and store all scratch files in.
@@ -455,14 +464,15 @@
         os.mkdir(self.FOREIGN_WORKING_TREE_PATH)
         return self.foreign_tree_store.fetch(self.FOREIGN_WORKING_TREE_PATH)
 
-    def importToBazaar(self, foreign_tree, bazaar_tree):
-        """Actually import `foreign_tree` into `bazaar_tree`.
+    def importToBazaar(self, foreign_tree, bazaar_branch):
+        """Actually import `foreign_tree` into `bazaar_branch`.
 
         :param foreign_tree: A `SubversionWorkingTree` or a `CVSWorkingTree`.
-        :param bazaar_tree: A `bzrlib.workingtree.WorkingTree`.
+        :param bazaar_tree: A `bzrlib.branch.Branch`, which must have a
+            colocated working tree.
         """
         foreign_directory = foreign_tree.local_path
-        bzr_directory = str(bazaar_tree.basedir)
+        bzr_directory = str(bazaar_branch.bzrdir.open_workingtree().basedir)
 
         scm_branch = SCM.branch(bzr_directory)
         last_commit = cscvs.findLastCscvsCommit(scm_branch)
@@ -500,9 +510,9 @@
 
     def _doImport(self):
         foreign_tree = self.getForeignTree()
-        bazaar_tree = self.getBazaarWorkingTree()
-        self.importToBazaar(foreign_tree, bazaar_tree)
-        non_trivial = self.pushBazaarWorkingTree(bazaar_tree)
+        bazaar_branch = self.getBazaarBranch()
+        self.importToBazaar(foreign_tree, bazaar_branch)
+        non_trivial = self.pushBazaarBranch(bazaar_branch)
         self.foreign_tree_store.archive(foreign_tree)
         if non_trivial:
             return CodeImportWorkerExitCode.SUCCESS
@@ -516,6 +526,8 @@
     Subclasses need to implement `format_classes`.
     """
 
+    needs_bzr_tree = False
+
     @property
     def format_classes(self):
         """The format classes that should be tried for this import."""
@@ -531,13 +543,14 @@
         return {}
 
     def _doImport(self):
-        bazaar_tree = self.getBazaarWorkingTree()
-        self.bazaar_branch_store.push(
-            self.source_details.branch_id, bazaar_tree, self.required_format)
+        self._logger.info("Starting job.")
         saved_factory = bzrlib.ui.ui_factory
         bzrlib.ui.ui_factory = LoggingUIFactory(
             writer=lambda m: self._logger.info('%s', m))
         try:
+            self._logger.info(
+                "Getting exising bzr branch from central store.")
+            bazaar_branch = self.getBazaarBranch()
             transport = get_transport(self.source_details.url)
             for format_class in self.format_classes:
                 try:
@@ -549,11 +562,14 @@
                 raise NotBranchError(self.source_details.url)
             foreign_branch = format.open(transport).open_branch()
             foreign_branch_tip = foreign_branch.last_revision()
-            inter_branch = InterBranch.get(foreign_branch, bazaar_tree.branch)
+            inter_branch = InterBranch.get(foreign_branch, bazaar_branch)
+            self._logger.info("Importing foreign branch.")
             pull_result = inter_branch.pull(
                 overwrite=True, **self.getExtraPullArgs())
-            self.pushBazaarWorkingTree(bazaar_tree)
-            last_imported_revison = bazaar_tree.branch.last_revision()
+            self._logger.info("Pushing foreign branch to central store.")
+            self.pushBazaarBranch(bazaar_branch)
+            last_imported_revison = bazaar_branch.last_revision()
+            self._logger.info("Job complete.")
             if last_imported_revison == foreign_branch_tip:
                 if pull_result.old_revid != pull_result.new_revid:
                     return CodeImportWorkerExitCode.SUCCESS
@@ -583,39 +599,38 @@
         """See `PullingImportWorker.getExtraPullArgs`."""
         return {'limit': config.codeimport.git_revisions_import_limit}
 
-    def getBazaarWorkingTree(self):
-        """See `ImportWorker.getBazaarWorkingTree`.
+    def getBazaarBranch(self):
+        """See `ImportWorker.getBazaarBranch`.
 
-        In addition to the superclass' behaviour, we retrieve the 'git.db'
-        shamap from the import data store and put it where bzr-git will find
-        it in the Bazaar tree, that is at '.bzr/repository/git.db'.
+        In addition to the superclass' behaviour, we retrieve bzr-git's
+        caches, both legacy and modern, from the import data store and put
+        them where bzr-git will find them in the Bazaar tree, that is at
+        '.bzr/repository/git.db' and '.bzr/repository/git'.
         """
-        tree = PullingImportWorker.getBazaarWorkingTree(self)
+        branch = PullingImportWorker.getBazaarBranch(self)
         # Fetch the legacy cache from the store, if present.
         self.import_data_store.fetch(
-            'git.db', tree.branch.repository._transport)
+            'git.db', branch.repository._transport)
         # The cache dir from newer bzr-gits is stored as a tarball.
         local_name = 'git-cache.tar.gz'
         if self.import_data_store.fetch(local_name):
-            repo_transport = tree.branch.repository._transport
+            repo_transport = branch.repository._transport
             repo_transport.mkdir('git')
             git_db_dir = os.path.join(
                 local_path_from_url(repo_transport.base), 'git')
             extract_tarball(local_name, git_db_dir)
-        return tree
-
-    def pushBazaarWorkingTree(self, bazaar_tree):
-        """See `ImportWorker.pushBazaarWorkingTree`.
-
-        In addition to the superclass' behaviour, we store the 'git.db' shamap
-        that bzr-git will have created at .bzr/repository/bzr.git into the
-        import data store.
+        return branch
+
+    def pushBazaarBranch(self, bazaar_branch):
+        """See `ImportWorker.pushBazaarBranch`.
+
+        In addition to the superclass' behaviour, we store bzr-git's cache
+        directory at .bzr/repository/git in the import data store.
         """
-        non_trivial = PullingImportWorker.pushBazaarWorkingTree(
-            self, bazaar_tree)
-        repo_transport = bazaar_tree.branch.repository._transport
-        git_db_dir = os.path.join(
-            local_path_from_url(repo_transport.base), 'git')
+        non_trivial = PullingImportWorker.pushBazaarBranch(
+            self, bazaar_branch)
+        repo_base = bazaar_branch.repository._transport.base
+        git_db_dir = os.path.join(local_path_from_url(repo_base), 'git')
         local_name = 'git-cache.tar.gz'
         create_tarball(git_db_dir, local_name)
         self.import_data_store.put(local_name)
@@ -637,29 +652,29 @@
         from bzrlib.plugins.hg import HgBzrDirFormat
         return [HgBzrDirFormat]
 
-    def getBazaarWorkingTree(self):
-        """See `ImportWorker.getBazaarWorkingTree`.
+    def getBazaarBranch(self):
+        """See `ImportWorker.getBazaarBranch`.
 
         In addition to the superclass' behaviour, we retrieve the 'hg-v2.db'
         map from the import data store and put it where bzr-hg will find
         it in the Bazaar tree, that is at '.bzr/repository/hg-v2.db'.
         """
-        tree = PullingImportWorker.getBazaarWorkingTree(self)
+        branch = PullingImportWorker.getBazaarBranch(self)
         self.import_data_store.fetch(
-            self.db_file, tree.branch.repository._transport)
-        return tree
-
-    def pushBazaarWorkingTree(self, bazaar_tree):
-        """See `ImportWorker.pushBazaarWorkingTree`.
-
-        In addition to the superclass' behaviour, we store the 'hg-v2.db' shamap
-        that bzr-hg will have created at .bzr/repository/hg-v2.db into the
-        import data store.
+            self.db_file, branch.repository._transport)
+        return branch
+
+    def pushBazaarBranch(self, bazaar_branch):
+        """See `ImportWorker.pushBazaarBranch`.
+
+        In addition to the superclass' behaviour, we store the 'hg-v2.db'
+        shamap that bzr-hg will have created at .bzr/repository/hg-v2.db into
+        the import data store.
         """
-        non_trivial = PullingImportWorker.pushBazaarWorkingTree(
-            self, bazaar_tree)
+        non_trivial = PullingImportWorker.pushBazaarBranch(
+            self, bazaar_branch)
         self.import_data_store.put(
-            self.db_file, bazaar_tree.branch.repository._transport)
+            self.db_file, bazaar_branch.repository._transport)
         return non_trivial