1
"""
2
Models for the REST interface
3
"""
4 4
import functools
5 4
import re
6 4
import warnings
7 4
from typing import Any, Dict, List, Optional, Tuple, Union
8

9 4
from pydantic import Field, constr, root_validator, validator
10 4
from qcelemental.util import get_base_docs
11

12 4
from .common_models import KeywordSet, Molecule, ObjectId, ProtoModel, KVStore
13 4
from .gridoptimization import GridOptimizationInput
14 4
from .records import ResultRecord
15 4
from .task_models import PriorityEnum, TaskRecord
16 4
from .torsiondrive import TorsionDriveInput
17

18 4
__all__ = [
19
    "ComputeResponse",
20
    "rest_model",
21
    "QueryStr",
22
    "QueryObjectId",
23
    "QueryListStr",
24
    "ResultResponse",
25
    "CollectionSubresourceGETResponseMeta",
26
]
27

28
### Utility functions
29

30 4
__rest_models = {}
31

32

33 4
def register_model(name: str, rest: str, body: ProtoModel, response: ProtoModel) -> None:
34
    """
35
    Registers a new REST model.
36

37
    Parameters
38
    ----------
39
    name : str
40
        A regular expression describing the rest endpoint.
41
    rest : str
42
        The REST endpoint type.
43
    body : ProtoModel
44
        The REST query body model.
45
    response : ProtoModel
46
        The REST query response model.
47

48
    """
49 4
    rest = rest.upper()
50

51 4
    if (name in __rest_models) and (rest in __rest_models[name]):
52 0
        raise KeyError(f"Model name {name} already registered.")
53

54 4
    if name not in __rest_models:
55 4
        __rest_models[name] = {}
56

57 4
    __rest_models[name][rest] = (body, response)
58

59

60 4
@functools.lru_cache(1000, typed=True)
61 4
def rest_model(resource: str, rest: str) -> Tuple[ProtoModel, ProtoModel]:
62
    """
63
    Acquires a REST Model.
64

65
    Parameters
66
    ----------
67
    resource : str
68
        The REST endpoint resource name.
69
    rest : str
70
        The REST endpoint type: GET, POST, PUT, DELETE
71

72
    Returns
73
    -------
74
    Tuple[ProtoModel, ProtoModel]
75
        The (body, response) models of the REST request.
76

77
    """
78 4
    rest = rest.upper()
79 4
    matches = []
80 4
    for model_re in __rest_models.keys():
81 4
        if re.fullmatch(model_re, resource):
82 4
            try:
83 4
                matches.append(__rest_models[model_re][rest])
84 0
            except KeyError:
85 0
                pass  # Could have different regexes for different endpoint types
86

87 4
    if len(matches) == 0:
88 0
        raise KeyError(f"REST Model for endpoint {resource} could not be found.")
89

90 4
    if len(matches) > 1:
91 0
        warnings.warn(
92
            f"Multiple REST models were matched for {rest} request at endpoint {resource}. "
93
            f"The following models will be used: {matches[0][0]}, {matches[0][1]}.",
94
            RuntimeWarning,
95
        )
96

97 4
    return matches[0]
98

99

100
### Generic Types and Common Models
101

102 4
nullstr = constr(regex="null")
103

104 4
QueryStr = Optional[Union[List[str], str]]
105 4
QueryInt = Optional[Union[List[int], int]]
106 4
QueryObjectId = Optional[Union[List[ObjectId], ObjectId]]
107 4
QueryNullObjectId = Optional[Union[List[ObjectId], ObjectId, List[nullstr], nullstr]]
108 4
QueryListStr = Optional[List[str]]
109

110

111 4
class EmptyMeta(ProtoModel):
112
    """
113
    There is no metadata accepted, so an empty metadata is sent for completion.
114
    """
115

116

117 4
class ResponseMeta(ProtoModel):
118
    """
119
    Standard Fractal Server response metadata
120
    """
121

122 4
    errors: List[Tuple[str, str]] = Field(
123
        ..., description="A list of error pairs in the form of [(error type, error message), ...]"
124
    )
125 4
    success: bool = Field(
126
        ...,
127
        description="Indicates if the passed information was successful in its duties. This is contextual to the "
128
        "data being passed in.",
129
    )
130 4
    error_description: Union[str, bool] = Field(
131
        ...,
132
        description="Details about the error if ``success`` is ``False``, otherwise this is ``False`` in the event "
133
        "of no errors.",
134
    )
135

136

137 4
class ResponseGETMeta(ResponseMeta):
138
    """
139
    Standard Fractal Server response metadata for GET/fetch type requests.
140
    """
141

142 4
    missing: List[str] = Field(..., description="The Id's of the objects which were not found in the database.")
143 4
    n_found: int = Field(
144
        ...,
145
        description="The number of entries which were already found in the database from the set which was provided.",
146
    )
147

148

149 4
class ResponsePOSTMeta(ResponseMeta):
150
    """
151
    Standard Fractal Server response metadata for POST/add type requests.
152
    """
153

154 4
    n_inserted: int = Field(
155
        ...,
156
        description="The number of new objects amongst the inputs which did not exist already, and are now in the "
157
        "database.",
158
    )
159 4
    duplicates: Union[List[str], List[Tuple[str, str]]] = Field(
160
        ...,
161
        description="The Ids of the objects which already exist in the database amongst the set which were passed in.",
162
    )
163 4
    validation_errors: List[str] = Field(
164
        ..., description="All errors with validating submitted objects will be documented here."
165
    )
166

167

168 4
class QueryMeta(ProtoModel):
169
    """
170
    Standard Fractal Server metadata for Database queries containing pagination information
171
    """
172

173 4
    limit: Optional[int] = Field(
174
        None, description="Limit to the number of objects which can be returned with this query."
175
    )
176 4
    skip: int = Field(0, description="The number of records to skip on the query.")
177

178

179 4
class QueryFilter(ProtoModel):
180
    """
181
    Standard Fractal Server metadata for column filtering
182
    """
183

184 4
    include: QueryListStr = Field(
185
        None,
186
        description="Return only these columns. Expert-level object. Only one of include and exclude may be specified.",
187
    )
188 4
    exclude: QueryListStr = Field(
189
        None,
190
        description="Return all but these columns. Expert-level object. Only one of include and exclude may be specified.",
191
    )
192

193 4
    @root_validator
194
    def check_include_or_exclude(cls, values):
195 4
        include = values.get("include")
196 4
        exclude = values.get("exclude")
197 4
        if (include is not None) and (exclude is not None):
198 4
            raise ValueError("Only one of include and exclude may be specified.")
199 4
        return values
200

201

202 4
class QueryMetaFilter(QueryMeta, QueryFilter):
203
    """
204
    Fractal Server metadata for Database queries allowing for filtering and pagination
205
    """
206

207

208 4
class ComputeResponse(ProtoModel):
209
    """
210
    The response model from the Fractal Server when new Compute or Services are added.
211
    """
212

213 4
    ids: List[Optional[ObjectId]] = Field(..., description="The Id's of the records to be computed.")
214 4
    submitted: List[ObjectId] = Field(
215
        ..., description="The object Ids which were submitted as new entries to the database."
216
    )
217 4
    existing: List[ObjectId] = Field(..., description="The list of object Ids which already existed in the database.")
218

219 4
    def __str__(self) -> str:
220 2
        return f"ComputeResponse(nsubmitted={len(self.submitted)} nexisting={len(self.existing)})"
221

222 4
    def __repr__(self) -> str:
223 0
        return f"<{self}>"
224

225 4
    def merge(self, other: "ComputeResponse") -> "ComputeResponse":
226
        """Merges two ComputeResponse objects together. The first takes precedence and order is maintained.
227

228
        Parameters
229
        ----------
230
        other : ComputeResponse
231
            The compute response to merge
232

233
        Returns
234
        -------
235
        ComputeResponse
236
            The merged compute response
237
        """
238 1
        return ComputeResponse(
239
            ids=(self.ids + other.ids),
240
            submitted=(self.submitted + other.submitted),
241
            existing=(self.existing + other.existing),
242
        )
243

244

245 4
common_docs = {
246
    EmptyMeta: str(get_base_docs(EmptyMeta)),
247
    ResponseMeta: str(get_base_docs(ResponseMeta)),
248
    ResponseGETMeta: str(get_base_docs(ResponseGETMeta)),
249
    ResponsePOSTMeta: str(get_base_docs(ResponsePOSTMeta)),
250
    QueryMeta: str(get_base_docs(QueryMeta)),
251
    QueryMetaFilter: str(get_base_docs(QueryMetaFilter)),
252
    ComputeResponse: str(get_base_docs(ComputeResponse)),
253
}
254

255
### Information
256

257

258 4
class InformationGETBody(ProtoModel):
259 4
    pass
260

261

262 4
class InformationGETResponse(ProtoModel):
263 4
    class Config(ProtoModel.Config):
264 4
        extra = "allow"
265

266

267 4
register_model("information", "GET", InformationGETBody, InformationGETResponse)
268

269
### KVStore
270

271

272 4
class KVStoreGETBody(ProtoModel):
273 4
    class Data(ProtoModel):
274 4
        id: QueryObjectId = Field(None, description="Id of the Key/Value Storage object to get.")
275

276 4
    meta: EmptyMeta = Field({}, description=common_docs[EmptyMeta])
277 4
    data: Data = Field(..., description="Data of the KV Get field: consists of Id of the Key/Value object to fetch.")
278

279

280 4
class KVStoreGETResponse(ProtoModel):
281 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
282 4
    data: Dict[str, KVStore] = Field(..., description="The entries of Key/Value object requested.")
283

284

285 4
register_model("kvstore", "GET", KVStoreGETBody, KVStoreGETResponse)
286

287
### Molecule response
288

289

290 4
class MoleculeGETBody(ProtoModel):
291 4
    class Data(ProtoModel):
292 4
        id: QueryObjectId = Field(None, description="Exact Id of the Molecule to fetch from the database.")
293 4
        molecule_hash: QueryStr = Field(
294
            None,
295
            description="Hash of the Molecule to search for in the database. Can be computed from the Molecule object "
296
            "directly without direct access to the Database itself.",
297
        )
298 4
        molecular_formula: QueryStr = Field(
299
            None,
300
            description="Query is made based on simple molecular formula. This is based on just the formula itself and "
301
            "contains no connectivity information.",
302
        )
303

304 4
    meta: QueryMeta = Field(QueryMeta(), description=common_docs[QueryMeta])
305 4
    data: Data = Field(
306
        ...,
307
        description="Data fields for a Molecule query.",  # Because Data is internal, this may not document sufficiently
308
    )
309

310

311 4
class MoleculeGETResponse(ProtoModel):
312 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
313 4
    data: List[Molecule] = Field(..., description="The List of Molecule objects found by the query.")
314

315

316 4
register_model("molecule", "GET", MoleculeGETBody, MoleculeGETResponse)
317

318

319 4
class MoleculePOSTBody(ProtoModel):
320 4
    meta: EmptyMeta = Field({}, description=common_docs[EmptyMeta])
321 4
    data: List[Molecule] = Field(..., description="A list of :class:`Molecule` objects to add to the Database.")
322

323

324 4
class MoleculePOSTResponse(ProtoModel):
325 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
326 4
    data: List[ObjectId] = Field(
327
        ...,
328
        description="A list of Id's assigned to the Molecule objects passed in which serves as a unique identifier "
329
        "in the database. If the Molecule was already in the database, then the Id returned is its "
330
        "existing Id (entries are not duplicated).",
331
    )
332

333

334 4
register_model("molecule", "POST", MoleculePOSTBody, MoleculePOSTResponse)
335

336
### Keywords
337

338

339 4
class KeywordGETBody(ProtoModel):
340 4
    class Data(ProtoModel):
341 4
        id: QueryObjectId = None
342 4
        hash_index: QueryStr = None
343

344 4
    meta: QueryMeta = Field(QueryMeta(), description=common_docs[QueryMeta])
345 4
    data: Data = Field(
346
        ...,
347
        description="The formal query for a Keyword fetch, contains ``id`` or ``hash_index`` for the object to fetch.",
348
    )
349

350

351 4
class KeywordGETResponse(ProtoModel):
352 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
353 4
    data: List[KeywordSet] = Field(
354
        ..., description="The :class:`KeywordSet` found from in the database based on the query."
355
    )
356

357

358 4
register_model("keyword", "GET", KeywordGETBody, KeywordGETResponse)
359

360

361 4
class KeywordPOSTBody(ProtoModel):
362 4
    meta: EmptyMeta = Field(
363
        {}, description="There is no metadata with this, so an empty metadata is sent for completion."
364
    )
365 4
    data: List[KeywordSet] = Field(..., description="The list of :class:`KeywordSet` objects to add to the database.")
366

367

368 4
class KeywordPOSTResponse(ProtoModel):
369 4
    data: List[Optional[ObjectId]] = Field(
370
        ...,
371
        description="The Ids assigned to the added :class:`KeywordSet` objects. In the event of duplicates, the Id "
372
        "will be the one already found in the database.",
373
    )
374 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
375

376

377 4
register_model("keyword", "POST", KeywordPOSTBody, KeywordPOSTResponse)
378

379
### Collections
380

381

382 4
class CollectionGETBody(ProtoModel):
383 4
    class Data(ProtoModel):
384 4
        collection: str = Field(
385
            None, description="The specific collection to look up as its identified in the database."
386
        )
387 4
        name: str = Field(None, description="The common name of the collection to look up.")
388

389 4
        @validator("collection")
390
        def cast_to_lower(cls, v):
391 4
            if v:
392 4
                v = v.lower()
393 4
            return v
394

395 4
    meta: QueryFilter = Field(
396
        None,
397
        description="Additional metadata to make with the query. Collections can only have an ``include/exclude`` key in its "
398
        "meta and therefore does not follow the standard GET metadata model.",
399
    )
400 4
    data: Data = Field(..., description="Information about the Collection to search the database with.")
401

402

403 4
class CollectionGETResponse(ProtoModel):
404 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
405 4
    data: List[Dict[str, Optional[Any]]] = Field(
406
        ..., description="The Collection objects returned by the server based on the query."
407
    )
408

409 4
    @validator("data")
410
    def ensure_collection_name_in_data_get_res(cls, v):
411 4
        for col in v:
412 4
            if "name" not in col or "collection" not in col:
413 0
                raise ValueError("Dicts in 'data' must have both 'collection' and 'name'")
414 4
        return v
415

416

417 4
register_model("collection", "GET", CollectionGETBody, CollectionGETResponse)
418

419

420 4
class CollectionPOSTBody(ProtoModel):
421 4
    class Meta(ProtoModel):
422 4
        overwrite: bool = Field(
423
            False,
424
            description="The existing Collection in the database will be updated if this is True, otherwise will "
425
            "remain unmodified if it already exists.",
426
        )
427

428 4
    class Data(ProtoModel):
429 4
        id: str = Field(
430
            "local",  # Auto blocks overwriting in a socket
431
            description="The Id of the object to assign in the database. If 'local', then it will not overwrite "
432
            "existing keys. There should be very little reason to ever touch this.",
433
        )
434 4
        collection: str = Field(
435
            ..., description="The specific identifier for this Collection as it will appear in database."
436
        )
437 4
        name: str = Field(..., description="The common name of this Collection.")
438

439 4
        class Config(ProtoModel.Config):
440 4
            extra = "allow"
441

442 4
        @validator("collection")
443
        def cast_to_lower(cls, v):
444 4
            return v.lower()
445

446 4
    meta: Meta = Field(
447
        Meta(),
448
        description="Metadata to specify how the Database should handle adding this Collection if it already exists. "
449
        "Metadata model for adding Collections can only accept ``overwrite`` as a key to choose to update "
450
        "existing Collections or not.",
451
    )
452 4
    data: Data = Field(..., description="The data associated with this Collection to add to the database.")
453

454

455 4
class CollectionPOSTResponse(ProtoModel):
456 4
    data: Union[str, None] = Field(
457
        ...,
458
        description="The Id of the Collection uniquely pointing to it in the Database. If the Collection was not added "
459
        "(e.g. ``overwrite=False`` for existing Collection), then a None is returned.",
460
    )
461 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
462

463

464 4
register_model("collection", "POST", CollectionPOSTBody, CollectionPOSTResponse)
465

466

467 4
class CollectionDELETEBody(ProtoModel):
468 4
    meta: EmptyMeta
469

470

471 4
class CollectionDELETEResponse(ProtoModel):
472 4
    meta: ResponseMeta
473

474

475 4
register_model("collection/[0-9]+", "DELETE", CollectionDELETEBody, CollectionDELETEResponse)
476

477
### Collection views
478

479

480 4
class CollectionSubresourceGETResponseMeta(ResponseMeta):
481
    """
482
    Response metadata for collection views functions.
483
    """
484

485 4
    msgpacked_cols: List[str] = Field(..., description="Names of columns which were serialized to msgpack-ext.")
486

487

488 4
class CollectionEntryGETBody(ProtoModel):
489 4
    class Data(ProtoModel):
490 4
        subset: QueryStr = Field(
491
            None,
492
            description="Not implemented. " "See qcfractal.interface.collections.dataset_view.DatasetView.get_entries",
493
        )
494

495 4
    meta: EmptyMeta = Field(EmptyMeta(), description=common_docs[EmptyMeta])
496 4
    data: Data = Field(..., description="Information about which entries to return.")
497

498

499 4
class CollectionEntryGETResponse(ProtoModel):
500 4
    meta: CollectionSubresourceGETResponseMeta = Field(
501
        ..., description=str(get_base_docs(CollectionSubresourceGETResponseMeta))
502
    )
503 4
    data: Optional[bytes] = Field(..., description="Feather-serialized bytes representing a pandas DataFrame.")
504

505

506 4
register_model("collection/[0-9]+/entry", "GET", CollectionEntryGETBody, CollectionEntryGETResponse)
507

508

509 4
class CollectionMoleculeGETBody(ProtoModel):
510 4
    class Data(ProtoModel):
511 4
        indexes: List[int] = Field(
512
            None,
513
            description="List of molecule indexes to return (returned by get_entries). "
514
            "See qcfractal.interface.collections.dataset_view.DatasetView.get_molecules",
515
        )
516

517 4
    meta: EmptyMeta = Field(EmptyMeta(), description=common_docs[EmptyMeta])
518 4
    data: Data = Field(..., description="Information about which molecules to return.")
519

520

521 4
class CollectionMoleculeGETResponse(ProtoModel):
522 4
    meta: CollectionSubresourceGETResponseMeta = Field(
523
        ..., description=str(get_base_docs(CollectionSubresourceGETResponseMeta))
524
    )
525 4
    data: Optional[bytes] = Field(..., description="Feather-serialized bytes representing a pandas DataFrame.")
526

527

528 4
register_model("collection/[0-9]+/molecule", "GET", CollectionMoleculeGETBody, CollectionMoleculeGETResponse)
529

530

531 4
class CollectionValueGETBody(ProtoModel):
532 4
    class Data(ProtoModel):
533 4
        class QueryData(ProtoModel):
534 4
            name: str
535 4
            driver: str
536 4
            native: bool
537

538 4
        queries: List[QueryData] = Field(
539
            None,
540
            description="List of queries to match against values columns. "
541
            "See qcfractal.interface.collections.dataset_view.DatasetView.get_values",
542
        )
543 4
        subset: QueryStr
544

545 4
    meta: EmptyMeta = Field(EmptyMeta(), description=common_docs[EmptyMeta])
546 4
    data: Data = Field(..., description="Information about which values to return.")
547

548

549 4
class CollectionValueGETResponse(ProtoModel):
550 4
    class Data(ProtoModel):
551 4
        values: bytes = Field(..., description="Feather-serialized bytes representing a pandas DataFrame.")
552 4
        units: Dict[str, str] = Field(..., description="Units of value columns.")
553

554 4
    meta: CollectionSubresourceGETResponseMeta = Field(
555
        ..., description=str(get_base_docs(CollectionSubresourceGETResponseMeta))
556
    )
557 4
    data: Optional[Data] = Field(..., description="Values and units.")
558

559

560 4
register_model("collection/[0-9]+/value", "GET", CollectionValueGETBody, CollectionValueGETResponse)
561

562

563 4
class CollectionListGETBody(ProtoModel):
564 4
    class Data(ProtoModel):
565 4
        pass
566

567 4
    meta: EmptyMeta = Field(EmptyMeta(), description=common_docs[EmptyMeta])
568 4
    data: Data = Field(..., description="Empty for now.")
569

570

571 4
class CollectionListGETResponse(ProtoModel):
572 4
    meta: CollectionSubresourceGETResponseMeta = Field(
573
        ..., description=str(get_base_docs(CollectionSubresourceGETResponseMeta))
574
    )
575 4
    data: Optional[bytes] = Field(..., description="Feather-serialized bytes representing a pandas DataFrame.")
576

577

578 4
register_model("collection/[0-9]+/list", "GET", CollectionListGETBody, CollectionListGETResponse)
579

580
### Result
581

582

583 4
class ResultGETBody(ProtoModel):
584 4
    class Data(ProtoModel):
585 4
        id: QueryObjectId = Field(
586
            None,
587
            description="The exact Id to fetch from the database. If this is set as a search condition, there is no "
588
            "reason to set anything else as this will be unique in the database, if it exists.",
589
        )
590 4
        task_id: QueryObjectId = Field(
591
            None,
592
            description="The exact Id of the task which carried out this Result's computation. If this is set as a "
593
            "search condition, there is no reason to set anything else as this will be unique in the "
594
            "database, if it exists. See also :class:`TaskRecord`.",
595
        )
596

597 4
        program: QueryStr = Field(
598
            None,
599
            description="Results will be searched to match the quantum chemistry software which carried out the "
600
            "calculation.",
601
        )
602 4
        molecule: QueryObjectId = Field(
603
            None, description="Results will be searched to match the Molecule Id which was computed on."
604
        )
605 4
        driver: QueryStr = Field(
606
            None,
607
            description="Results will be searched to match what class of computation was done. "
608
            "See :class:`DriverEnum` for valid choices and more information.",
609
        )
610 4
        method: QueryStr = Field(
611
            None,
612
            description="Results will be searched to match the quantum chemistry method executed to compute the value.",
613
        )
614 4
        basis: QueryStr = Field(
615
            None,
616
            description="Results will be searched to match specified basis sets which were used to compute the values.",
617
        )
618 4
        keywords: QueryNullObjectId = Field(
619
            None,
620
            description="Results will be searched based on which :class:`KeywordSet` was used to run the computation.",
621
        )
622

623 4
        status: QueryStr = Field(
624
            "COMPLETE",
625
            description="Results will be searched based on where they are in the compute pipeline. See the "
626
            ":class:`RecordStatusEnum` for valid statuses and more information.",
627
        )
628

629 4
        @validator("keywords", each_item=True, pre=True)
630
        def validate_keywords(cls, v):
631 4
            if v is None:
632 0
                v = "null"
633 4
            return v
634

635 4
        @validator("basis", each_item=True, pre=True)
636
        def validate_basis(cls, v):
637 4
            if (v is None) or (v == ""):
638 2
                v = "null"
639 4
            return v
640

641 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
642 4
    data: Data = Field(
643
        ..., description="The keys with data to search the database on for individual quantum chemistry computations."
644
    )
645

646

647 4
class ResultGETResponse(ProtoModel):
648 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
649
    # Either a record or dict depending if projection
650 4
    data: Union[List[ResultRecord], List[Dict[str, Any]]] = Field(
651
        ...,
652
        description="Results found from the query. This is a list of :class:`ResultRecord` in most cases, however, "
653
        "if a projection was specified in the GET request, then a dict is returned with mappings based "
654
        "on the projection.",
655
    )
656

657 4
    @validator("data", pre=True)
658
    def ensure_list_of_dict(cls, v):
659 4
        if isinstance(v, dict):
660 0
            return [v]
661 4
        return v
662

663

664 4
register_model("result", "GET", ResultGETBody, ResultGETResponse)
665

666
### Wavefunction data
667

668

669 4
class WavefunctionStoreGETBody(ProtoModel):
670 4
    class Data(ProtoModel):
671 4
        id: ObjectId = Field(None, description="Id of the Wavefunction Key/Value Storage object to get.")
672

673 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
674 4
    data: Data = Field(
675
        ...,
676
        description="Data of the Wavefunction Get field: consists of a ObjectId of the Wavefunction object to fetch.",
677
    )
678

679

680 4
class WavefunctionStoreGETResponse(ProtoModel):
681 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
682 4
    data: Dict[str, Any] = Field(..., description="The entries of the Wavefunction object requested.")
683

684

685 4
register_model("wavefunctionstore", "GET", WavefunctionStoreGETBody, WavefunctionStoreGETResponse)
686

687
### Procedures
688

689

690 4
class ProcedureGETBody(ProtoModel):
691 4
    class Data(ProtoModel):
692 4
        id: QueryObjectId = Field(
693
            None,
694
            description="The exact Id to fetch from the database. If this is set as a search condition, there is no "
695
            "reason to set anything else as this will be unique in the database, if it exists.",
696
        )
697 4
        task_id: QueryObjectId = Field(
698
            None,
699
            description="The exact Id of a task which is carried out by this Procedure. If this is set as a "
700
            "search condition, there is no reason to set anything else as this will be unique in the "
701
            "database, if it exists. See also :class:`TaskRecord`.",
702
        )
703

704 4
        procedure: QueryStr = Field(None, description="Procedures will be searched based on the name of the procedure.")
705 4
        program: QueryStr = Field(
706
            None,
707
            description="Procedures will be searched based on the program which is the main manager of the procedure",
708
        )
709 4
        hash_index: QueryStr = Field(
710
            None,
711
            description="Procedures will be searched based on a hash of the defined procedure. This is something which "
712
            "can be generated by the Procedure spec itself and does not require server access to compute. "
713
            "This should be unique in the database so there should be no reason to set anything else "
714
            "if this is set as a query.",
715
        )
716 4
        status: QueryStr = Field(
717
            "COMPLETE",
718
            description="Procedures will be searched based on where they are in the compute pipeline. See the "
719
            ":class:`RecordStatusEnum` for valid statuses.",
720
        )
721

722 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
723 4
    data: Data = Field(..., description="The keys with data to search the database on for Procedures.")
724

725

726 4
class ProcedureGETResponse(ProtoModel):
727 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
728 4
    data: List[Dict[str, Optional[Any]]] = Field(
729
        ..., description="The list of Procedure specs found based on the query."
730
    )
731

732

733 4
register_model("procedure", "GET", ProcedureGETBody, ProcedureGETResponse)
734

735
### Task Queue
736

737

738 4
class TaskQueueGETBody(ProtoModel):
739 4
    class Data(ProtoModel):
740 4
        id: QueryObjectId = Field(
741
            None,
742
            description="The exact Id to fetch from the database. If this is set as a search condition, there is no "
743
            "reason to set anything else as this will be unique in the database, if it exists.",
744
        )
745 4
        hash_index: QueryStr = Field(
746
            None,
747
            description="Tasks will be searched based on a hash of the defined Task. This is something which can "
748
            "be generated by the Task spec itself and does not require server access to compute. "
749
            "This should be unique in the database so there should be no reason to set anything else "
750
            "if this is set as a query.",
751
        )
752 4
        program: QueryStr = Field(
753
            None, description="Tasks will be searched based on the program responsible for executing this task."
754
        )
755 4
        status: QueryStr = Field(
756
            None,
757
            description="Tasks will be search based on where they are in the compute pipeline. See the "
758
            ":class:`RecordStatusEnum` for valid statuses.",
759
        )
760 4
        base_result: QueryStr = Field(
761
            None,
762
            description="The exact Id of the Result which this Task is linked to. If this is set as a "
763
            "search condition, there is no reason to set anything else as this will be unique in the "
764
            "database, if it exists. See also :class:`ResultRecord`.",
765
        )
766 4
        tag: QueryStr = Field(None, description="Tasks will be searched based on their associated tag.")
767 4
        manager: QueryStr = Field(
768
            None, description="Tasks will be searched based on the manager responsible for executing the task."
769
        )
770

771 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
772 4
    data: Data = Field(..., description="The keys with data to search the database on for Tasks.")
773

774

775 4
class TaskQueueGETResponse(ProtoModel):
776 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
777 4
    data: Union[List[TaskRecord], List[Dict[str, Any]]] = Field(
778
        ...,
779
        description="Tasks found from the query. This is a list of :class:`TaskRecord` in most cases, however, "
780
        "if a projection was specified in the GET request, then a dict is returned with mappings based "
781
        "on the projection.",
782
    )
783

784

785 4
register_model("task_queue", "GET", TaskQueueGETBody, TaskQueueGETResponse)
786

787

788 4
class TaskQueuePOSTBody(ProtoModel):
789 4
    class Meta(ProtoModel):
790 4
        procedure: str = Field(..., description="Name of the procedure which the Task will execute.")
791 4
        program: str = Field(..., description="The program which this Task will execute.")
792

793 4
        tag: Optional[str] = Field(
794
            None,
795
            description="Tag to assign to this Task so that Queue Managers can pull only Tasks based on this entry."
796
            "If no Tag is specified, any Queue Manager can pull this Task.",
797
        )
798 4
        priority: Union[PriorityEnum, None] = Field(None, description=str(PriorityEnum.__doc__))
799

800 4
        class Config(ProtoModel.Config):
801 4
            extra = "allow"
802

803 4
        @validator("priority", pre=True)
804
        def munge_priority(cls, v):
805 4
            if isinstance(v, str):
806 4
                v = PriorityEnum[v.upper()]
807 4
            return v
808

809 4
    meta: Meta = Field(..., description="The additional specification information for the Task to add to the Database.")
810 4
    data: List[Union[ObjectId, Molecule]] = Field(
811
        ...,
812
        description="The list of either Molecule objects or Molecule Id's (those already in the database) to submit as "
813
        "part of this Task.",
814
    )
815

816

817 4
class TaskQueuePOSTResponse(ProtoModel):
818

819 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
820 4
    data: ComputeResponse = Field(..., description="Data returned from the server from adding a Task.")
821

822

823 4
register_model("task_queue", "POST", TaskQueuePOSTBody, TaskQueuePOSTResponse)
824

825

826 4
class TaskQueuePUTBody(ProtoModel):
827 4
    class Data(ProtoModel):
828 4
        id: QueryObjectId = Field(
829
            None,
830
            description="The exact Id to target in database. If this is set as a search condition, there is no "
831
            "reason to set anything else as this will be unique in the database, if it exists.",
832
        )
833 4
        base_result: QueryObjectId = Field(  # TODO: Validate this description is correct
834
            None,
835
            description="The exact Id of a result which this Task is slated to write to. If this is set as a "
836
            "search condition, there is no reason to set anything else as this will be unique in the "
837
            "database, if it exists. See also :class:`ResultRecord`.",
838
        )
839

840 4
    class Meta(ProtoModel):
841 4
        operation: str = Field(..., description="The specific action you are taking as part of this update.")
842

843 4
        @validator("operation")
844
        def cast_to_lower(cls, v):
845 2
            return v.lower()
846

847 4
    meta: Meta = Field(..., description="The instructions to pass to the target Task from ``data``.")
848 4
    data: Data = Field(..., description="The information which contains the Task target in the database.")
849

850

851 4
class TaskQueuePUTResponse(ProtoModel):
852 4
    class Data(ProtoModel):
853 4
        n_updated: int = Field(..., description="The number of tasks which were changed.")
854

855 4
    meta: ResponseMeta = Field(..., description=common_docs[ResponseMeta])
856 4
    data: Data = Field(..., description="Information returned from attempting updates of Tasks.")
857

858

859 4
register_model("task_queue", "PUT", TaskQueuePUTBody, TaskQueuePUTResponse)
860

861
### Service Queue
862

863

864 4
class ServiceQueueGETBody(ProtoModel):
865 4
    class Data(ProtoModel):
866 4
        id: QueryObjectId = Field(
867
            None,
868
            description="The exact Id to fetch from the database. If this is set as a search condition, there is no "
869
            "reason to set anything else as this will be unique in the database, if it exists.",
870
        )
871 4
        procedure_id: QueryObjectId = Field(  # TODO: Validate this description is correct
872
            None,
873
            description="The exact Id of the Procedure this Service is responsible for executing. If this is set as a "
874
            "search condition, there is no reason to set anything else as this will be unique in the "
875
            "database, if it exists.",
876
        )
877 4
        hash_index: QueryStr = Field(
878
            None,
879
            description="Services are searched based on a hash of the defined Service. This is something which can "
880
            "be generated by the Service spec itself and does not require server access to compute. "
881
            "This should be unique in the database so there should be no reason to set anything else "
882
            "if this is set as a query.",
883
        )
884 4
        status: QueryStr = Field(
885
            None,
886
            description="Services are searched based on where they are in the compute pipeline. See the "
887
            ":class:`RecordStatusEnum` for valid statuses.",
888
        )
889

890 4
    meta: QueryMeta = Field(QueryMeta(), description=common_docs[QueryMeta])
891 4
    data: Data = Field(..., description="The keys with data to search the database on for Services.")
892

893

894 4
class ServiceQueueGETResponse(ProtoModel):
895 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
896 4
    data: List[Dict[str, Optional[Any]]] = Field(
897
        ..., description="The return of Services found in the database mapping their Ids to the Service spec."
898
    )
899

900

901 4
register_model("service_queue", "GET", ServiceQueueGETBody, ServiceQueueGETResponse)
902

903

904 4
class ServiceQueuePOSTBody(ProtoModel):
905 4
    class Meta(ProtoModel):
906 4
        tag: Optional[str] = Field(
907
            None,
908
            description="Tag to assign to the Tasks this Service will generate so that Queue Managers can pull only "
909
            "Tasks based on this entry. If no Tag is specified, any Queue Manager can pull this Tasks "
910
            "created by this Service.",
911
        )
912 4
        priority: Union[str, int, None] = Field(
913
            None,
914
            description="Priority given to this Tasks created by this Service. Higher priority will be pulled first.",
915
        )
916

917 4
    meta: Meta = Field(
918
        ...,
919
        description="Metadata information for the Service for the Tag and Priority of Tasks this Service will create.",
920
    )
921 4
    data: List[Union[TorsionDriveInput, GridOptimizationInput]] = Field(
922
        ..., description="A list the specification for Procedures this Service will manage and generate Tasks for."
923
    )
924

925

926 4
class ServiceQueuePOSTResponse(ProtoModel):
927

928 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
929 4
    data: ComputeResponse = Field(..., description="Data returned from the server from adding a Service.")
930

931

932 4
register_model("service_queue", "POST", ServiceQueuePOSTBody, ServiceQueuePOSTResponse)
933

934

935 4
class ServiceQueuePUTBody(ProtoModel):
936 4
    class Data(ProtoModel):
937 4
        id: QueryObjectId = Field(None, description="The Id of the Service.")
938 4
        procedure_id: QueryObjectId = Field(None, description="The Id of the Procedure that the Service is linked to.")
939

940 4
    class Meta(ProtoModel):
941 4
        operation: str = Field(..., description="The update action to perform.")
942

943 4
        @validator("operation")
944
        def cast_to_lower(cls, v):
945 1
            return v.lower()
946

947 4
    meta: Meta = Field(..., description="The instructions to pass to the targeted Service.")
948 4
    data: Data = Field(..., description="The information which contains the Service target in the database.")
949

950

951 4
class ServiceQueuePUTResponse(ProtoModel):
952 4
    class Data(ProtoModel):
953 4
        n_updated: int = Field(..., description="The number of services which were changed.")
954

955 4
    meta: ResponseMeta = Field(..., description=common_docs[ResponseMeta])
956 4
    data: Data = Field(..., description="Information returned from attempting updates of Services.")
957

958

959 4
register_model("service_queue", "PUT", ServiceQueuePUTBody, ServiceQueuePUTResponse)
960

961
### Queue Manager
962

963

964 4
class QueueManagerMeta(ProtoModel):
965
    """
966
    Validation and identification Meta information for the Queue Manager's communication with the Fractal Server.
967
    """
968

969
    # Name data
970 4
    cluster: str = Field(..., description="The Name of the Cluster the Queue Manager is running on.")
971 4
    hostname: str = Field(..., description="Hostname of the machine the Queue Manager is running on.")
972 4
    uuid: str = Field(..., description="A UUID assigned to the QueueManager to uniquely identify it.")
973

974
    # Username
975 4
    username: Optional[str] = Field(None, description="Fractal Username the Manager is being executed under.")
976

977
    # Version info
978 4
    qcengine_version: str = Field(..., description="Version of QCEngine which the Manager has access to.")
979 4
    manager_version: str = Field(
980
        ..., description="Version of the QueueManager (Fractal) which is getting and returning Jobs."
981
    )
982

983
    # search info
984 4
    programs: List[str] = Field(
985
        ...,
986
        description="A list of programs which the QueueManager, and thus QCEngine, has access to. Affects which Tasks "
987
        "the Manager can pull.",
988
    )
989 4
    procedures: List[str] = Field(
990
        ...,
991
        description="A list of procedures which the QueueManager has access to. Affects which Tasks "
992
        "the Manager can pull.",
993
    )
994 4
    tag: QueryStr = Field(
995
        None,
996
        description="Optional queue tag to pull Tasks from. If None, tasks are pulled from all tags. "
997
        "If a list of tags is provided, tasks are pulled in order of tags. (This does not "
998
        "guarantee tasks will be executed in that order, however.)",
999
    )
1000

1001
    # Statistics
1002 4
    total_worker_walltime: Optional[float] = Field(None, description="The total worker walltime in core-hours.")
1003 4
    total_task_walltime: Optional[float] = Field(None, description="The total task walltime in core-hours.")
1004 4
    active_tasks: Optional[int] = Field(None, description="The total number of active running tasks.")
1005 4
    active_cores: Optional[int] = Field(None, description="The total number of active cores.")
1006 4
    active_memory: Optional[float] = Field(None, description="The total amount of active memory in GB.")
1007

1008

1009
# Add the new QueueManagerMeta to the docs
1010 4
common_docs[QueueManagerMeta] = str(get_base_docs(QueueManagerMeta))
1011

1012

1013 4
class QueueManagerGETBody(ProtoModel):
1014 4
    class Data(ProtoModel):
1015 4
        limit: int = Field(..., description="Max number of Queue Managers to get from the server.")
1016

1017 4
    meta: QueueManagerMeta = Field(..., description=common_docs[QueueManagerMeta])
1018 4
    data: Data = Field(
1019
        ...,
1020
        description="A model of Task request data for the Queue Manager to fetch. Accepts ``limit`` as the maximum "
1021
        "number of tasks to pull.",
1022
    )
1023

1024

1025 4
class QueueManagerGETResponse(ProtoModel):
1026 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
1027 4
    data: List[Dict[str, Optional[Any]]] = Field(
1028
        ..., description="A list of tasks retrieved from the server to compute."
1029
    )
1030

1031

1032 4
register_model("queue_manager", "GET", QueueManagerGETBody, QueueManagerGETResponse)
1033

1034

1035 4
class QueueManagerPOSTBody(ProtoModel):
1036 4
    meta: QueueManagerMeta = Field(..., description=common_docs[QueueManagerMeta])
1037 4
    data: Dict[ObjectId, Any] = Field(..., description="A Dictionary of tasks to return to the server.")
1038

1039

1040 4
class QueueManagerPOSTResponse(ProtoModel):
1041 4
    meta: ResponsePOSTMeta = Field(..., description=common_docs[ResponsePOSTMeta])
1042 4
    data: bool = Field(..., description="A True/False return on if the server accepted the returned tasks.")
1043

1044

1045 4
register_model("queue_manager", "POST", QueueManagerPOSTBody, QueueManagerPOSTResponse)
1046

1047

1048 4
class QueueManagerPUTBody(ProtoModel):
1049 4
    class Data(ProtoModel):
1050 4
        operation: str
1051 4
        configuration: Optional[Dict[str, Any]] = None
1052

1053 4
    meta: QueueManagerMeta = Field(..., description=common_docs[QueueManagerMeta])
1054 4
    data: Data = Field(
1055
        ...,
1056
        description="The update action which the Queue Manager requests the Server take with respect to how the "
1057
        "Queue Manager is tracked.",
1058
    )
1059

1060

1061 4
class QueueManagerPUTResponse(ProtoModel):
1062 4
    meta: Dict[str, Any] = Field({}, description=common_docs[EmptyMeta])
1063
    # Order on Union[] is important. Union[bool, Dict[str, int]] -> True if the input dict is not empty since
1064
    # Python can resolve dict -> bool since it passes a `is` test. Will not cast bool -> dict[str, int], so make Dict[]
1065
    # check first
1066 4
    data: Union[Dict[str, int], bool] = Field(
1067
        ...,
1068
        description="The response from the Server attempting to update the Queue Manager's server-side status. "
1069
        "Response type is a function of the operation made from the PUT request.",
1070
    )
1071

1072

1073 4
register_model("queue_manager", "PUT", QueueManagerPUTBody, QueueManagerPUTResponse)
1074

1075
## advanced procedures queries
1076

1077

1078 4
class OptimizationFinalResultBody(ProtoModel):
1079 4
    class Data(ProtoModel):
1080 4
        optimization_ids: QueryObjectId = Field(
1081
            None, description="List of optimization procedure Ids to fetch their final results from the database."
1082
        )
1083

1084
    # TODO: not yet supported
1085 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
1086 4
    data: Data = Field(..., description="The keys with data to search the database on for Procedures.")
1087

1088

1089 4
class OptimizationAllResultBody(ProtoModel):
1090 4
    class Data(ProtoModel):
1091 4
        optimization_ids: QueryObjectId = Field(
1092
            None, description="List of optimization procedure Ids to fetch their ALL their results from the database."
1093
        )
1094

1095
    # TODO: not yet supported
1096 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
1097 4
    data: Data = Field(..., description="The keys with data to search the database on for Procedures.")
1098

1099

1100 4
class OptimizationInitialMoleculeBody(ProtoModel):
1101 4
    class Data(ProtoModel):
1102 4
        optimization_ids: QueryObjectId = Field(
1103
            None, description="List of optimization procedure Ids to fetch their initial  molecules from the database."
1104
        )
1105

1106
    # TODO: not yet supported
1107 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
1108 4
    data: Data = Field(..., description="The keys with data to search the database on for Procedures.")
1109

1110

1111 4
class OptimizationFinalMoleculeBody(ProtoModel):
1112 4
    class Data(ProtoModel):
1113 4
        optimization_ids: QueryObjectId = Field(
1114
            None, description="List of optimization procedure Ids to fetch their final molecules from the database."
1115
        )
1116

1117
    # TODO: not yet supported
1118 4
    meta: QueryMetaFilter = Field(QueryMetaFilter(), description=common_docs[QueryMetaFilter])
1119 4
    data: Data = Field(..., description="The keys with data to search the database on for Procedures.")
1120

1121

1122 4
class ResultResponse(ProtoModel):
1123 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
1124
    # Either a record or dict depending if projection
1125 4
    data: Union[Dict[str, ResultRecord], Dict[str, Any]] = Field(
1126
        ..., description="A List of Results found from the query per optimization id."
1127
    )
1128

1129

1130 4
class ListResultResponse(ProtoModel):
1131 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
1132
    # Either a record or dict depending if projection
1133 4
    data: Union[Dict[str, List[ResultRecord]], Dict[str, Any]] = Field(
1134
        ..., description="A List of Results found from the query per optimization id."
1135
    )
1136

1137

1138 4
class ListMoleculeResponse(ProtoModel):
1139 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
1140
    # Either a record or dict depending if projection
1141 4
    data: Union[Dict[str, Molecule], Dict[str, Any]] = Field(
1142
        ..., description="A List of Molecules found from the query per optimization id."
1143
    )
1144

1145

1146 4
register_model(r"optimization/final_result", "GET", OptimizationFinalResultBody, ResultResponse)
1147 4
register_model(r"optimization/all_results", "GET", OptimizationAllResultBody, ListResultResponse)
1148 4
register_model(r"optimization/initial_molecule", "GET", OptimizationAllResultBody, ListMoleculeResponse)
1149 4
register_model(r"optimization/final_molecule", "GET", OptimizationAllResultBody, ListMoleculeResponse)
1150

1151

1152 4
class ManagerInfoGETBody(ProtoModel):
1153 4
    class Data(ProtoModel):
1154 4
        name: QueryStr = Field(None, description="Name(s) of managers to query for.")
1155 4
        status: QueryStr = Field(
1156
            None,
1157
            description="Managers will be searched based on status. See :class:`ManagerStatusEnum` for valid statuses.",
1158
        )
1159

1160 4
    meta: QueryMeta = Field(QueryMeta(), description=common_docs[QueryMeta])
1161 4
    data: Data = Field(..., description="The keys with data to search the database on for Managers.")
1162

1163

1164 4
class ManagerInfoGETResponse(ProtoModel):
1165 4
    meta: ResponseGETMeta = Field(..., description=common_docs[ResponseGETMeta])
1166 4
    data: List[Dict[str, Any]] = Field(..., description="Information about the requested managers")
1167

1168

1169 4
register_model(r"manager", "GET", ManagerInfoGETBody, ManagerInfoGETResponse)

Read our documentation on viewing source code .

Loading