sa.py 3.32 KB
Newer Older
Dylan Jones's avatar
no slug    
Dylan Jones committed
1
2
from contextlib import contextmanager

3
4
from sqlalchemy import String, ForeignKey, Enum, Integer, Float, JSON
from sqlalchemy.ext.declarative import as_declarative
Dylan Jones's avatar
Dylan Jones committed
5
from sqlalchemy.orm import relationship, validates
Dylan Jones's avatar
Dylan Jones committed
6
from sqlalchemy.schema import Column
Dylan Jones's avatar
Dylan Jones committed
7
8

from .field_types import FieldType
Dylan Jones's avatar
Dylan Jones committed
9
from .meta import Session
Dylan Jones's avatar
Dylan Jones committed
10
11


Dylan Jones's avatar
no slug    
Dylan Jones committed
12
@contextmanager
Dylan Jones's avatar
Dylan Jones committed
13
def session_context():
Dylan Jones's avatar
no slug    
Dylan Jones committed
14
15
16
17
18
19
20
21
22
23
24
    session = Session()
    try:
        yield session
        session.commit()
    except BaseException:
        session.rollback()
        raise
    finally:
        session.close()


25
@as_declarative()
Dylan Jones's avatar
Dylan Jones committed
26
class Base(object):
Dylan Jones's avatar
Dylan Jones committed
27
    pass
Dylan Jones's avatar
Dylan Jones committed
28
29


Dylan Jones's avatar
Dylan Jones committed
30
31
class Point(Base):
    """
32
    Represents actual instances of any and all points on the map.
Dylan Jones's avatar
Dylan Jones committed
33
34
    """
    __tablename__ = 'point'
Dylan Jones's avatar
Dylan Jones committed
35
    id = Column(Integer, primary_key=True, autoincrement=True)
Dylan Jones's avatar
Dylan Jones committed
36
37
38
    name = Column(String, nullable=True)
    lat = Column(Float, nullable=False)
    lon = Column(Float, nullable=False)
39
    attributes = Column(JSON, nullable=False)
Dylan Jones's avatar
Dylan Jones committed
40
41
42
43
44

    # Relationships
    category_id = Column(Integer, ForeignKey('category.id'), nullable=False)
    category = relationship('Category')
    parent_id = Column(Integer, ForeignKey('point.id'), nullable=True)
Dylan Jones's avatar
Dylan Jones committed
45
46
    parent = relationship('Point', remote_side=[id])
    children = relationship('Point')
Dylan Jones's avatar
Dylan Jones committed
47
48

    @validates('attributes')
Dylan Jones's avatar
Dylan Jones committed
49
50
51
    def validate_data(self, _, data):
        if data is None:
            return
Dylan Jones's avatar
Dylan Jones committed
52
        fields = self.category.fields
Dylan Jones's avatar
Dylan Jones committed
53
        for key in data:
Dylan Jones's avatar
Dylan Jones committed
54
            # Find Field object that corresponds to this key
Dylan Jones's avatar
Dylan Jones committed
55
            for field in fields:
Dylan Jones's avatar
Dylan Jones committed
56
                if field.slug == key:
Dylan Jones's avatar
Dylan Jones committed
57
58
                    break
            else:
Dylan Jones's avatar
Dylan Jones committed
59
                raise ValueError(f'extra data "{key}" must be a registered field')
Dylan Jones's avatar
Dylan Jones committed
60
61
62
            field.validate_data(data[key])
        return data

Dylan Jones's avatar
Dylan Jones committed
63
64
65
66
67
68
69
70
71
72
73
74
    def as_json(self, children=True):
        if children:
            children = [child.as_json(children=False) for child in self.children]
        return {
            "name": self.name,
            "lat": self.lat,
            "lon": self.lon,
            "category": self.category.id,
            "attributes": self.attributes,
            "children": children
        }

Dylan Jones's avatar
Dylan Jones committed
75
76

class Category(Base):
Dylan Jones's avatar
Dylan Jones committed
77
    """
78
    Represent a schema for a single category of objects (e.g. water fountain or bathroom)
Dylan Jones's avatar
Dylan Jones committed
79
    """
Dylan Jones's avatar
Dylan Jones committed
80
    __tablename__ = 'category'
Dylan Jones's avatar
Dylan Jones committed
81

Dylan Jones's avatar
Dylan Jones committed
82
    id = Column(Integer, primary_key=True, autoincrement=True)
Dylan Jones's avatar
Dylan Jones committed
83
    name = Column(String, nullable=False, unique=True)
Dylan Jones's avatar
Dylan Jones committed
84
    icon = Column(String, nullable=True)
Dylan Jones's avatar
Dylan Jones committed
85
86

    fields = relationship("Field")
Dylan Jones's avatar
Dylan Jones committed
87

Dylan Jones's avatar
Dylan Jones committed
88
89
90
91
92
93
94
95
    def as_json(self):
        return {
            "id": self.id,
            "name": self.name,
            "icon": self.icon,
            "attributes": {attr.slug: attr.as_json() for attr in self.fields}
        }

Dylan Jones's avatar
Dylan Jones committed
96
97
98

class Field(Base):
    """
99
    Represents a single field in the Category schema.
Dylan Jones's avatar
Dylan Jones committed
100
    """
Dylan Jones's avatar
Dylan Jones committed
101
    __tablename__ = 'field'
Dylan Jones's avatar
Dylan Jones committed
102

Dylan Jones's avatar
Dylan Jones committed
103
    id = Column(Integer, primary_key=True, autoincrement=True)
Dylan Jones's avatar
Dylan Jones committed
104
    slug = Column(String, nullable=False)
Dylan Jones's avatar
Dylan Jones committed
105
    name = Column(String, nullable=False)
Dylan Jones's avatar
Dylan Jones committed
106
    type = Column(Enum(FieldType), nullable=False)
Dylan Jones's avatar
Dylan Jones committed
107
108
109

    # Relationship
    category_id = Column(Integer, ForeignKey('category.id'))
Dylan Jones's avatar
Dylan Jones committed
110
111
112
113
114

    def validate_data(self, data):
        """
        Verify that data is the correct type for this Field.
        """
115
        self.type.validate(data)
Dylan Jones's avatar
Dylan Jones committed
116
117
118
119
120
121
122

    def as_json(self):
        return {
            "slug": self.slug,
            "name": self.name,
            "type": self.type.name
        }