Allow "bin/drydock lease ..." to select particular blueprints with "--blueprint"

Summary: Ref T13676. See discussion in that task.

Test Plan: Tried nonsense blueprint identifiers, duplicate identifiers, incompatible identifiers, and valid identifiers. Got apparently sensible behavior in all cases.

Subscribers: PHID-OPKG-gm6ozazyms6q6i22gyam

Maniphest Tasks: T13676

Differential Revision: https://secure.phabricator.com/D21802

src/applications/drydock/management/DrydockManagementLeaseWorkflow.php

@@ -32,6 +32,12 @@ final class DrydockManagementLeaseWorkflow
             'default' => 1,
             'help' => pht('Lease a given number of identical resources.'),
           ),
+          array(
+            'name' => 'blueprint',
+            'param' => 'identifier',
+            'repeat' => true,
+            'help' => pht('Lease resources from a specific blueprint.'),
+          ),
         ));
   }
 
@@ -79,6 +85,13 @@ final class DrydockManagementLeaseWorkflow
       $attributes = array();
     }
 
+    $filter_identifiers = $args->getArg('blueprint');
+    if ($filter_identifiers) {
+      $filter_blueprints = $this->getBlueprintFilterMap($filter_identifiers);
+    } else {
+      $filter_blueprints = array();
+    }
+
     $blueprint_phids = null;
 
     $leases = array();
@@ -94,7 +107,9 @@ final class DrydockManagementLeaseWorkflow
       }
 
       if ($blueprint_phids === null) {
-        $blueprint_phids = $this->newAllowedBlueprintPHIDs($lease);
+        $blueprint_phids = $this->newAllowedBlueprintPHIDs(
+          $lease,
+          $filter_blueprints);
       }
 
       $lease->setAllowedBlueprintPHIDs($blueprint_phids);
@@ -253,7 +268,51 @@ final class DrydockManagementLeaseWorkflow
     }
   }
 
-  private function newAllowedBlueprintPHIDs(DrydockLease $lease) {
+  private function getBlueprintFilterMap(array $identifiers) {
+    $viewer = $this->getViewer();
+
+    $query = id(new DrydockBlueprintQuery())
+      ->setViewer($viewer)
+      ->withIdentifiers($identifiers);
+
+    $blueprints = $query->execute();
+    $blueprints = mpull($blueprints, null, 'getPHID');
+
+    $map = $query->getIdentifierMap();
+
+    $seen = array();
+    foreach ($identifiers as $identifier) {
+      if (!isset($map[$identifier])) {
+        throw new PhutilArgumentUsageException(
+          pht(
+            'Blueprint "%s" could not be loaded. Try a blueprint ID or '.
+            'PHID.',
+            $identifier));
+      }
+
+      $blueprint = $map[$identifier];
+
+      $blueprint_phid = $blueprint->getPHID();
+      if (isset($seen[$blueprint_phid])) {
+        throw new PhutilArgumentUsageException(
+          pht(
+            'Blueprint "%s" is specified more than once (as "%s" and "%s").',
+            $blueprint->getBlueprintName(),
+            $seen[$blueprint_phid],
+            $identifier));
+      }
+
+      $seen[$blueprint_phid] = true;
+    }
+
+    return mpull($map, null, 'getPHID');
+  }
+
+  private function newAllowedBlueprintPHIDs(
+    DrydockLease $lease,
+    array $filter_blueprints) {
+    assert_instances_of($filter_blueprints, 'DrydockBlueprint');
+
     $viewer = $this->getViewer();
 
     $impls = DrydockBlueprintImplementation::getAllForAllocatingLease($lease);
@@ -282,6 +341,23 @@ final class DrydockManagementLeaseWorkflow
 
     $phids = mpull($blueprints, 'getPHID');
 
+    if ($filter_blueprints) {
+      $allowed_map = array_fuse($phids);
+      $filter_map = mpull($filter_blueprints, null, 'getPHID');
+
+      foreach ($filter_map as $filter_phid => $blueprint) {
+        if (!isset($allowed_map[$filter_phid])) {
+          throw new PhutilArgumentUsageException(
+            pht(
+              'Specified blueprint "%s" is not capable of satisfying the '.
+              'configured lease.',
+              $blueprint->getBlueprintName()));
+        }
+      }
+
+      $phids = mpull($filter_blueprints, 'getPHID');
+    }
+
     return $phids;
   }
 

src/applications/drydock/query/DrydockBlueprintQuery.php

@@ -9,6 +9,11 @@ final class DrydockBlueprintQuery extends DrydockQuery {
   private $disabled;
   private $authorizedPHIDs;
 
+  private $identifiers;
+  private $identifierIDs;
+  private $identifierPHIDs;
+  private $identifierMap;
+
   public function withIDs(array $ids) {
     $this->ids = $ids;
     return $this;
@@ -45,6 +50,43 @@ final class DrydockBlueprintQuery extends DrydockQuery {
       $ngrams);
   }
 
+  public function withIdentifiers(array $identifiers) {
+    if (!$identifiers) {
+      throw new Exception(
+        pht(
+          'Can not issue a query with an empty identifier list.'));
+    }
+
+    $this->identifiers = $identifiers;
+
+    $ids = array();
+    $phids = array();
+
+    foreach ($identifiers as $identifier) {
+      if (ctype_digit($identifier)) {
+        $ids[] = $identifier;
+      } else {
+        $phids[] = $identifier;
+      }
+    }
+
+    $this->identifierIDs = $ids;
+    $this->identifierPHIDs = $phids;
+
+    return $this;
+  }
+
+  public function getIdentifierMap() {
+    if ($this->identifierMap === null) {
+      throw new Exception(
+        pht(
+          'Execute a query with identifiers before getting the '.
+          'identifier map.'));
+    }
+
+    return $this->identifierMap;
+  }
+
   public function newResultObject() {
     return new DrydockBlueprint();
   }
@@ -57,6 +99,14 @@ final class DrydockBlueprintQuery extends DrydockQuery {
     return $this->loadStandardPage($this->newResultObject());
   }
 
+  protected function willExecute() {
+    if ($this->identifiers) {
+      $this->identifierMap = array();
+    } else {
+      $this->identifierMap = null;
+    }
+  }
+
   protected function willFilterPage(array $blueprints) {
     $impls = DrydockBlueprintImplementation::getAllBlueprintImplementations();
     foreach ($blueprints as $key => $blueprint) {
@@ -70,6 +120,30 @@ final class DrydockBlueprintQuery extends DrydockQuery {
       $blueprint->attachImplementation($impl);
     }
 
+    if ($this->identifiers) {
+      $id_map = mpull($blueprints, null, 'getID');
+      $phid_map = mpull($blueprints, null, 'getPHID');
+
+      $map = $this->identifierMap;
+
+      foreach ($this->identifierIDs as $id) {
+        if (isset($id_map[$id])) {
+          $map[$id] = $id_map[$id];
+        }
+      }
+
+      foreach ($this->identifierPHIDs as $phid) {
+        if (isset($phid_map[$phid])) {
+          $map[$phid] = $phid_map[$phid];
+        }
+      }
+
+      // Just for consistency, reorder the map to match input order.
+      $map = array_select_keys($map, $this->identifiers);
+
+      $this->identifierMap = $map;
+    }
+
     return $blueprints;
   }
 
@@ -111,6 +185,29 @@ final class DrydockBlueprintQuery extends DrydockQuery {
         (int)$this->disabled);
     }
 
+    if ($this->identifiers !== null) {
+      $parts = array();
+
+      if ($this->identifierIDs) {
+        $parts[] = qsprintf(
+          $conn,
+          'blueprint.id IN (%Ld)',
+          $this->identifierIDs);
+      }
+
+      if ($this->identifierPHIDs) {
+        $parts[] = qsprintf(
+          $conn,
+          'blueprint.phid IN (%Ls)',
+          $this->identifierPHIDs);
+      }
+
+      $where[] = qsprintf(
+        $conn,
+        '%LO',
+        $parts);
+    }
+
     return $where;
   }