diff --git a/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php b/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php --- a/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php +++ b/src/applications/drydock/management/DrydockManagementLeaseWorkflow.php @@ -32,6 +32,12 @@ '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 @@ $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 @@ } if ($blueprint_phids === null) { - $blueprint_phids = $this->newAllowedBlueprintPHIDs($lease); + $blueprint_phids = $this->newAllowedBlueprintPHIDs( + $lease, + $filter_blueprints); } $lease->setAllowedBlueprintPHIDs($blueprint_phids); @@ -253,7 +268,51 @@ } } - 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 @@ $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; } diff --git a/src/applications/drydock/query/DrydockBlueprintQuery.php b/src/applications/drydock/query/DrydockBlueprintQuery.php --- a/src/applications/drydock/query/DrydockBlueprintQuery.php +++ b/src/applications/drydock/query/DrydockBlueprintQuery.php @@ -9,6 +9,11 @@ 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 @@ $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 @@ 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 @@ $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 @@ (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; }