1 4
from sqlalchemy import JSON, Boolean, Column, ForeignKey, Index, Integer, String
2 4
from sqlalchemy.dialects import postgresql
3 4
from sqlalchemy.ext.hybrid import hybrid_property
4 4
from sqlalchemy.orm import relationship
5 4
from sqlalchemy.sql.functions import GenericFunction
6

7 4
from qcfractal.storage_sockets.models.sql_base import Base, MsgpackExt
8

9
# class json_agg(GenericFunction):
10
#     type = postgresql.JSON
11

12

13 4
class json_build_object(GenericFunction):
14 4
    type = postgresql.JSON
15

16

17 4
class CollectionORM(Base):
18
    """
19
    A base collection class of precomuted workflows such as datasets, ..
20

21
    This is a dynamic document, so it will accept any number of
22
    extra fields (expandable and uncontrolled schema)
23
    """
24

25 4
    __tablename__ = "collection"
26

27 4
    id = Column(Integer, primary_key=True)
28 4
    collection_type = Column(String)  # for inheritance
29

30 4
    collection = Column(String(100), nullable=False)
31 4
    lname = Column(String(100), nullable=False)
32 4
    name = Column(String(100), nullable=False)
33

34 4
    tags = Column(JSON)
35 4
    tagline = Column(String)
36 4
    description = Column(String)
37

38 4
    group = Column(String(100), nullable=False)
39 4
    visibility = Column(Boolean, nullable=False)
40

41 4
    view_url_hdf5 = Column(String)
42 4
    view_url_plaintext = Column(String)
43 4
    view_metadata = Column(JSON)
44 4
    view_available = Column(Boolean, nullable=False)
45

46 4
    provenance = Column(JSON)
47

48 4
    extra = Column(JSON)  # extra data related to specific collection type
49

50 4
    def update_relations(self, **kwarg):
51 4
        pass
52

53 4
    __table_args__ = (
54
        Index("ix_collection_lname", "collection", "lname", unique=True),
55
        Index("ix_collection_type", "collection_type"),
56
    )
57

58 4
    __mapper_args__ = {"polymorphic_on": "collection_type"}
59

60

61
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
62

63

64 4
class DatasetMixin:
65
    """
66
    Mixin class for common Dataset attributes.
67
    """
68

69 4
    default_benchmark = Column(String)
70 4
    default_keywords = Column(JSON)
71

72 4
    default_driver = Column(String)
73 4
    default_units = Column(String)
74 4
    alias_keywords = Column(JSON)
75 4
    default_program = Column(String)
76

77 4
    history_keys = Column(JSON)
78 4
    history = Column(JSON)
79

80

81 4
class ContributedValuesORM(Base):
82
    """One group of a contibuted values per dataset
83
    Each dataset can have multiple rows in this table"""
84

85 4
    __tablename__ = "contributed_values"
86

87 4
    collection_id = Column(Integer, ForeignKey("collection.id", ondelete="cascade"), primary_key=True)
88

89 4
    name = Column(String, nullable=False, primary_key=True)
90 4
    values = Column(MsgpackExt, nullable=False)
91 4
    index = Column(MsgpackExt, nullable=False)
92 4
    values_structure = Column(JSON, nullable=False)
93

94 4
    theory_level = Column(JSON, nullable=False)
95 4
    units = Column(String, nullable=False)
96 4
    theory_level_details = Column(JSON)
97

98 4
    citations = Column(JSON)
99 4
    external_url = Column(String)
100 4
    doi = Column(String)
101

102 4
    comments = Column(String)
103

104

105 4
class DatasetEntryORM(Base):
106
    """Association table for many to many"""
107

108 4
    __tablename__ = "dataset_entry"
109

110 4
    dataset_id = Column(Integer, ForeignKey("dataset.id", ondelete="cascade"), primary_key=True)
111
    # TODO: check the cascase_delete with molecule
112 4
    molecule_id = Column(Integer, ForeignKey("molecule.id"), nullable=False)
113

114 4
    name = Column(String, nullable=False, primary_key=True)
115 4
    comment = Column(String)
116 4
    local_results = Column(JSON)
117

118

119 4
class DatasetORM(CollectionORM, DatasetMixin):
120
    """
121
    The Dataset class for homogeneous computations on many molecules.
122
    """
123

124 4
    __tablename__ = "dataset"
125

126 4
    id = Column(Integer, ForeignKey("collection.id", ondelete="CASCADE"), primary_key=True)
127

128 4
    contributed_values_obj = relationship(ContributedValuesORM, lazy="selectin", cascade="all, delete-orphan")
129

130 4
    records_obj = relationship(
131
        DatasetEntryORM, lazy="selectin", cascade="all, delete-orphan", backref="dataset"  # lazy='noload',
132
    )
133

134 4
    @hybrid_property
135
    def contributed_values(self):
136 4
        return self._contributed_values(self.contributed_values_obj)
137

138 4
    @staticmethod
139
    def _contributed_values(contributed_values_obj):
140 4
        if not contributed_values_obj:
141 4
            return {}
142

143 4
        if not isinstance(contributed_values_obj, list):
144 4
            contributed_values_obj = [contributed_values_obj]
145 4
        ret = {}
146 4
        try:
147 4
            for obj in contributed_values_obj:
148 4
                ret[obj.name.lower()] = obj.to_dict(exclude=["collection_id"])
149 4
        except Exception as err:
150 4
            pass
151

152 4
        return ret
153

154 4
    @contributed_values.setter
155
    def contributed_values(self, dict_values):
156 4
        return dict_values
157

158 4
    @hybrid_property
159
    def records(self):
160
        """calculated property when accessed, not saved in the DB
161
        A view of the many to many relation"""
162

163 4
        return self._records(self.records_obj)
164

165 4
    @staticmethod
166
    def _records(records_obj):
167

168 4
        if not records_obj:
169 4
            return []
170

171 4
        if not isinstance(records_obj, list):
172 4
            records_obj = [records_obj]
173

174 4
        ret = []
175 4
        try:
176 4
            for rec in records_obj:
177 4
                ret.append(rec.to_dict(exclude=["dataset_id"]))
178 4
        except Exception as err:
179
            # raises exception of first access!!
180 4
            pass
181

182 4
        return ret
183

184 4
    @records.setter
185
    def records(self, dict_values):
186 4
        return dict_values
187

188 4
    def update_relations(self, records=None, contributed_values=None, **kwarg):
189

190 4
        self.records_obj = []
191 4
        records = [] if not records else records
192 4
        for rec_dict in records:
193 4
            rec = DatasetEntryORM(dataset_id=int(self.id), **rec_dict)
194 4
            self.records_obj.append(rec)
195

196 4
        self.contributed_values_obj = []
197 4
        contributed_values = {} if not contributed_values else contributed_values
198 4
        for key, rec_dict in contributed_values.items():
199 4
            rec = ContributedValuesORM(collection_id=int(self.id), **rec_dict)
200 4
            self.contributed_values_obj.append(rec)
201

202 4
    __table_args__ = (
203
        # Index('ix_results_molecule', 'molecule'),  # b-tree index
204
        # UniqueConstraint("program", "driver", "method", "basis", "keywords", "molecule", name='uix_results_keys'),
205
    )
206

207 4
    __mapper_args__ = {
208
        "polymorphic_identity": "dataset",
209
        # to have separate select when querying CollectionORM
210
        "polymorphic_load": "selectin",
211
    }
212

213

214
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
215

216

217 4
class ReactionDatasetEntryORM(Base):
218
    """Association table for many to many"""
219

220 4
    __tablename__ = "reaction_dataset_entry"
221

222 4
    reaction_dataset_id = Column(Integer, ForeignKey("reaction_dataset.id", ondelete="cascade"), primary_key=True)
223

224 4
    attributes = Column(JSON)
225 4
    name = Column(String, nullable=False, primary_key=True)
226 4
    reaction_results = Column(JSON)
227 4
    stoichiometry = Column(JSON)
228 4
    extras = Column(JSON)
229

230

231 4
class ReactionDatasetORM(CollectionORM, DatasetMixin):
232
    """
233
    Reaction Dataset
234
    """
235

236 4
    __tablename__ = "reaction_dataset"
237

238 4
    id = Column(Integer, ForeignKey("collection.id", ondelete="CASCADE"), primary_key=True)
239

240 4
    ds_type = Column(String)
241

242 4
    records_obj = relationship(
243
        ReactionDatasetEntryORM, lazy="selectin", cascade="all, delete-orphan", backref="reaction_dataset"
244
    )
245

246 4
    contributed_values_obj = relationship(ContributedValuesORM, lazy="selectin", cascade="all, delete-orphan")
247

248 4
    @hybrid_property
249
    def contributed_values(self):
250 4
        return self._contributed_values(self.contributed_values_obj)
251

252 4
    @staticmethod
253
    def _contributed_values(contributed_values_obj):
254 4
        return DatasetORM._contributed_values(contributed_values_obj)
255

256 4
    @contributed_values.setter
257
    def contributed_values(self, dict_values):
258 4
        return dict_values
259

260 4
    def update_relations(self, records=None, contributed_values=None, **kwarg):
261

262 4
        self.records_obj = []
263 4
        records = records or []
264 4
        for rec_dict in records:
265 4
            rec = ReactionDatasetEntryORM(reaction_dataset_id=int(self.id), **rec_dict)
266 4
            self.records_obj.append(rec)
267

268 4
        self.contributed_values_obj = []
269 4
        contributed_values = {} if not contributed_values else contributed_values
270 4
        for key, rec_dict in contributed_values.items():
271 4
            rec = ContributedValuesORM(collection_id=int(self.id), **rec_dict)
272 4
            self.contributed_values_obj.append(rec)
273

274 4
    @hybrid_property
275
    def records(self):
276
        """calculated property when accessed, not saved in the DB
277
        A view of the many to many relation"""
278

279 4
        return self._records(self.records_obj)
280

281 4
    @staticmethod
282
    def _records(records_obj):
283

284 4
        if not records_obj:
285 4
            return []
286

287 4
        if not isinstance(records_obj, list):
288 4
            records_obj = [records_obj]
289

290 4
        ret = []
291 4
        try:
292 4
            for rec in records_obj:
293 4
                ret.append(rec.to_dict(exclude=["reaction_dataset_id"]))
294 4
        except Exception as err:
295
            # raises exception of first access!!
296 4
            pass
297 4
        return ret
298

299 4
    @records.setter
300
    def records(self, dict_values):
301 4
        return dict_values
302

303 4
    __table_args__ = (
304
        # Index('ix_results_molecule', 'molecule'),  # b-tree index
305
        # UniqueConstraint("program", "driver", "method", "basis", "keywords", "molecule", name='uix_results_keys'),
306
    )
307

308 4
    __mapper_args__ = {
309
        "polymorphic_identity": "reactiondataset",
310
        # to have separate select when querying CollectionORM
311
        "polymorphic_load": "selectin",
312
    }
313

314

315
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Read our documentation on viewing source code .

Loading