Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from sqlalchemy import schema as sa_schema 

2from sqlalchemy.types import Integer 

3from sqlalchemy.types import NULLTYPE 

4 

5from .. import util 

6from ..util.compat import string_types 

7 

8 

9class SchemaObjects(object): 

10 def __init__(self, migration_context=None): 

11 self.migration_context = migration_context 

12 

13 def primary_key_constraint(self, name, table_name, cols, schema=None): 

14 m = self.metadata() 

15 columns = [sa_schema.Column(n, NULLTYPE) for n in cols] 

16 t = sa_schema.Table(table_name, m, *columns, schema=schema) 

17 p = sa_schema.PrimaryKeyConstraint(*[t.c[n] for n in cols], name=name) 

18 t.append_constraint(p) 

19 return p 

20 

21 def foreign_key_constraint( 

22 self, 

23 name, 

24 source, 

25 referent, 

26 local_cols, 

27 remote_cols, 

28 onupdate=None, 

29 ondelete=None, 

30 deferrable=None, 

31 source_schema=None, 

32 referent_schema=None, 

33 initially=None, 

34 match=None, 

35 **dialect_kw 

36 ): 

37 m = self.metadata() 

38 if source == referent and source_schema == referent_schema: 

39 t1_cols = local_cols + remote_cols 

40 else: 

41 t1_cols = local_cols 

42 sa_schema.Table( 

43 referent, 

44 m, 

45 *[sa_schema.Column(n, NULLTYPE) for n in remote_cols], 

46 schema=referent_schema 

47 ) 

48 

49 t1 = sa_schema.Table( 

50 source, 

51 m, 

52 *[sa_schema.Column(n, NULLTYPE) for n in t1_cols], 

53 schema=source_schema 

54 ) 

55 

56 tname = ( 

57 "%s.%s" % (referent_schema, referent) 

58 if referent_schema 

59 else referent 

60 ) 

61 

62 dialect_kw["match"] = match 

63 

64 f = sa_schema.ForeignKeyConstraint( 

65 local_cols, 

66 ["%s.%s" % (tname, n) for n in remote_cols], 

67 name=name, 

68 onupdate=onupdate, 

69 ondelete=ondelete, 

70 deferrable=deferrable, 

71 initially=initially, 

72 **dialect_kw 

73 ) 

74 t1.append_constraint(f) 

75 

76 return f 

77 

78 def unique_constraint(self, name, source, local_cols, schema=None, **kw): 

79 t = sa_schema.Table( 

80 source, 

81 self.metadata(), 

82 *[sa_schema.Column(n, NULLTYPE) for n in local_cols], 

83 schema=schema 

84 ) 

85 kw["name"] = name 

86 uq = sa_schema.UniqueConstraint(*[t.c[n] for n in local_cols], **kw) 

87 # TODO: need event tests to ensure the event 

88 # is fired off here 

89 t.append_constraint(uq) 

90 return uq 

91 

92 def check_constraint(self, name, source, condition, schema=None, **kw): 

93 t = sa_schema.Table( 

94 source, 

95 self.metadata(), 

96 sa_schema.Column("x", Integer), 

97 schema=schema, 

98 ) 

99 ck = sa_schema.CheckConstraint(condition, name=name, **kw) 

100 t.append_constraint(ck) 

101 return ck 

102 

103 def generic_constraint(self, name, table_name, type_, schema=None, **kw): 

104 t = self.table(table_name, schema=schema) 

105 types = { 

106 "foreignkey": lambda name: sa_schema.ForeignKeyConstraint( 

107 [], [], name=name 

108 ), 

109 "primary": sa_schema.PrimaryKeyConstraint, 

110 "unique": sa_schema.UniqueConstraint, 

111 "check": lambda name: sa_schema.CheckConstraint("", name=name), 

112 None: sa_schema.Constraint, 

113 } 

114 try: 

115 const = types[type_] 

116 except KeyError: 

117 raise TypeError( 

118 "'type' can be one of %s" 

119 % ", ".join(sorted(repr(x) for x in types)) 

120 ) 

121 else: 

122 const = const(name=name) 

123 t.append_constraint(const) 

124 return const 

125 

126 def metadata(self): 

127 kw = {} 

128 if ( 

129 self.migration_context is not None 

130 and "target_metadata" in self.migration_context.opts 

131 ): 

132 mt = self.migration_context.opts["target_metadata"] 

133 if hasattr(mt, "naming_convention"): 

134 kw["naming_convention"] = mt.naming_convention 

135 return sa_schema.MetaData(**kw) 

136 

137 def table(self, name, *columns, **kw): 

138 m = self.metadata() 

139 t = sa_schema.Table(name, m, *columns, **kw) 

140 for f in t.foreign_keys: 

141 self._ensure_table_for_fk(m, f) 

142 return t 

143 

144 def column(self, name, type_, **kw): 

145 return sa_schema.Column(name, type_, **kw) 

146 

147 def index(self, name, tablename, columns, schema=None, **kw): 

148 t = sa_schema.Table( 

149 tablename or "no_table", self.metadata(), schema=schema 

150 ) 

151 idx = sa_schema.Index( 

152 name, 

153 *[util.sqla_compat._textual_index_column(t, n) for n in columns], 

154 **kw 

155 ) 

156 return idx 

157 

158 def _parse_table_key(self, table_key): 

159 if "." in table_key: 

160 tokens = table_key.split(".") 

161 sname = ".".join(tokens[0:-1]) 

162 tname = tokens[-1] 

163 else: 

164 tname = table_key 

165 sname = None 

166 return (sname, tname) 

167 

168 def _ensure_table_for_fk(self, metadata, fk): 

169 """create a placeholder Table object for the referent of a 

170 ForeignKey. 

171 

172 """ 

173 if isinstance(fk._colspec, string_types): 

174 table_key, cname = fk._colspec.rsplit(".", 1) 

175 sname, tname = self._parse_table_key(table_key) 

176 if table_key not in metadata.tables: 

177 rel_t = sa_schema.Table(tname, metadata, schema=sname) 

178 else: 

179 rel_t = metadata.tables[table_key] 

180 if cname not in rel_t.c: 

181 rel_t.append_column(sa_schema.Column(cname, NULLTYPE))