introspection.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  1. from collections import namedtuple
  2. from django.db.backends.base.introspection import BaseDatabaseIntrospection
  3. from django.db.backends.base.introspection import FieldInfo as BaseFieldInfo
  4. from django.db.backends.base.introspection import TableInfo as BaseTableInfo
  5. from django.db.models import Index
  6. FieldInfo = namedtuple("FieldInfo", BaseFieldInfo._fields + ("is_autofield", "comment"))
  7. TableInfo = namedtuple("TableInfo", BaseTableInfo._fields + ("comment",))
  8. class DatabaseIntrospection(BaseDatabaseIntrospection):
  9. # Maps type codes to Django Field types.
  10. data_types_reverse = {
  11. 16: "BooleanField",
  12. 17: "BinaryField",
  13. 20: "BigIntegerField",
  14. 21: "SmallIntegerField",
  15. 23: "IntegerField",
  16. 25: "TextField",
  17. 700: "FloatField",
  18. 701: "FloatField",
  19. 869: "GenericIPAddressField",
  20. 1042: "CharField", # blank-padded
  21. 1043: "CharField",
  22. 1082: "DateField",
  23. 1083: "TimeField",
  24. 1114: "DateTimeField",
  25. 1184: "DateTimeField",
  26. 1186: "DurationField",
  27. 1266: "TimeField",
  28. 1700: "DecimalField",
  29. 2950: "UUIDField",
  30. 3802: "JSONField",
  31. }
  32. # A hook for subclasses.
  33. index_default_access_method = "btree"
  34. ignored_tables = []
  35. def get_field_type(self, data_type, description):
  36. field_type = super().get_field_type(data_type, description)
  37. if description.is_autofield or (
  38. # Required for pre-Django 4.1 serial columns.
  39. description.default
  40. and "nextval" in description.default
  41. ):
  42. if field_type == "IntegerField":
  43. return "AutoField"
  44. elif field_type == "BigIntegerField":
  45. return "BigAutoField"
  46. elif field_type == "SmallIntegerField":
  47. return "SmallAutoField"
  48. return field_type
  49. def get_table_list(self, cursor):
  50. """Return a list of table and view names in the current database."""
  51. cursor.execute(
  52. """
  53. SELECT
  54. c.relname,
  55. CASE
  56. WHEN c.relispartition THEN 'p'
  57. WHEN c.relkind IN ('m', 'v') THEN 'v'
  58. ELSE 't'
  59. END,
  60. obj_description(c.oid, 'pg_class')
  61. FROM pg_catalog.pg_class c
  62. LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
  63. WHERE c.relkind IN ('f', 'm', 'p', 'r', 'v')
  64. AND n.nspname NOT IN ('pg_catalog', 'pg_toast')
  65. AND pg_catalog.pg_table_is_visible(c.oid)
  66. """
  67. )
  68. return [
  69. TableInfo(*row)
  70. for row in cursor.fetchall()
  71. if row[0] not in self.ignored_tables
  72. ]
  73. def get_table_description(self, cursor, table_name):
  74. """
  75. Return a description of the table with the DB-API cursor.description
  76. interface.
  77. """
  78. # Query the pg_catalog tables as cursor.description does not reliably
  79. # return the nullable property and information_schema.columns does not
  80. # contain details of materialized views.
  81. cursor.execute(
  82. """
  83. SELECT
  84. a.attname AS column_name,
  85. NOT (a.attnotnull OR (t.typtype = 'd' AND t.typnotnull)) AS is_nullable,
  86. pg_get_expr(ad.adbin, ad.adrelid) AS column_default,
  87. CASE WHEN collname = 'default' THEN NULL ELSE collname END AS collation,
  88. a.attidentity != '' AS is_autofield,
  89. col_description(a.attrelid, a.attnum) AS column_comment
  90. FROM pg_attribute a
  91. LEFT JOIN pg_attrdef ad ON a.attrelid = ad.adrelid AND a.attnum = ad.adnum
  92. LEFT JOIN pg_collation co ON a.attcollation = co.oid
  93. JOIN pg_type t ON a.atttypid = t.oid
  94. JOIN pg_class c ON a.attrelid = c.oid
  95. JOIN pg_namespace n ON c.relnamespace = n.oid
  96. WHERE c.relkind IN ('f', 'm', 'p', 'r', 'v')
  97. AND c.relname = %s
  98. AND n.nspname NOT IN ('pg_catalog', 'pg_toast')
  99. AND pg_catalog.pg_table_is_visible(c.oid)
  100. """,
  101. [table_name],
  102. )
  103. field_map = {line[0]: line[1:] for line in cursor.fetchall()}
  104. cursor.execute(
  105. "SELECT * FROM %s LIMIT 1" % self.connection.ops.quote_name(table_name)
  106. )
  107. return [
  108. FieldInfo(
  109. line.name,
  110. line.type_code,
  111. # display_size is always None on psycopg2.
  112. line.internal_size if line.display_size is None else line.display_size,
  113. line.internal_size,
  114. line.precision,
  115. line.scale,
  116. *field_map[line.name],
  117. )
  118. for line in cursor.description
  119. ]
  120. def get_sequences(self, cursor, table_name, table_fields=()):
  121. cursor.execute(
  122. """
  123. SELECT
  124. s.relname AS sequence_name,
  125. a.attname AS colname
  126. FROM
  127. pg_class s
  128. JOIN pg_depend d ON d.objid = s.oid
  129. AND d.classid = 'pg_class'::regclass
  130. AND d.refclassid = 'pg_class'::regclass
  131. JOIN pg_attribute a ON d.refobjid = a.attrelid
  132. AND d.refobjsubid = a.attnum
  133. JOIN pg_class tbl ON tbl.oid = d.refobjid
  134. AND tbl.relname = %s
  135. AND pg_catalog.pg_table_is_visible(tbl.oid)
  136. WHERE
  137. s.relkind = 'S';
  138. """,
  139. [table_name],
  140. )
  141. return [
  142. {"name": row[0], "table": table_name, "column": row[1]}
  143. for row in cursor.fetchall()
  144. ]
  145. def get_relations(self, cursor, table_name):
  146. """
  147. Return a dictionary of {field_name: (field_name_other_table, other_table)}
  148. representing all foreign keys in the given table.
  149. """
  150. cursor.execute(
  151. """
  152. SELECT a1.attname, c2.relname, a2.attname
  153. FROM pg_constraint con
  154. LEFT JOIN pg_class c1 ON con.conrelid = c1.oid
  155. LEFT JOIN pg_class c2 ON con.confrelid = c2.oid
  156. LEFT JOIN
  157. pg_attribute a1 ON c1.oid = a1.attrelid AND a1.attnum = con.conkey[1]
  158. LEFT JOIN
  159. pg_attribute a2 ON c2.oid = a2.attrelid AND a2.attnum = con.confkey[1]
  160. WHERE
  161. c1.relname = %s AND
  162. con.contype = 'f' AND
  163. c1.relnamespace = c2.relnamespace AND
  164. pg_catalog.pg_table_is_visible(c1.oid)
  165. """,
  166. [table_name],
  167. )
  168. return {row[0]: (row[2], row[1]) for row in cursor.fetchall()}
  169. def get_constraints(self, cursor, table_name):
  170. """
  171. Retrieve any constraints or keys (unique, pk, fk, check, index) across
  172. one or more columns. Also retrieve the definition of expression-based
  173. indexes.
  174. """
  175. constraints = {}
  176. # Loop over the key table, collecting things as constraints. The column
  177. # array must return column names in the same order in which they were
  178. # created.
  179. cursor.execute(
  180. """
  181. SELECT
  182. c.conname,
  183. array(
  184. SELECT attname
  185. FROM unnest(c.conkey) WITH ORDINALITY cols(colid, arridx)
  186. JOIN pg_attribute AS ca ON cols.colid = ca.attnum
  187. WHERE ca.attrelid = c.conrelid
  188. ORDER BY cols.arridx
  189. ),
  190. c.contype,
  191. (SELECT fkc.relname || '.' || fka.attname
  192. FROM pg_attribute AS fka
  193. JOIN pg_class AS fkc ON fka.attrelid = fkc.oid
  194. WHERE fka.attrelid = c.confrelid AND fka.attnum = c.confkey[1]),
  195. cl.reloptions
  196. FROM pg_constraint AS c
  197. JOIN pg_class AS cl ON c.conrelid = cl.oid
  198. WHERE cl.relname = %s AND pg_catalog.pg_table_is_visible(cl.oid)
  199. """,
  200. [table_name],
  201. )
  202. for constraint, columns, kind, used_cols, options in cursor.fetchall():
  203. constraints[constraint] = {
  204. "columns": columns,
  205. "primary_key": kind == "p",
  206. "unique": kind in ["p", "u"],
  207. "foreign_key": tuple(used_cols.split(".", 1)) if kind == "f" else None,
  208. "check": kind == "c",
  209. "index": False,
  210. "definition": None,
  211. "options": options,
  212. }
  213. # Now get indexes
  214. cursor.execute(
  215. """
  216. SELECT
  217. indexname,
  218. array_agg(attname ORDER BY arridx),
  219. indisunique,
  220. indisprimary,
  221. array_agg(ordering ORDER BY arridx),
  222. amname,
  223. exprdef,
  224. s2.attoptions
  225. FROM (
  226. SELECT
  227. c2.relname as indexname, idx.*, attr.attname, am.amname,
  228. CASE
  229. WHEN idx.indexprs IS NOT NULL THEN
  230. pg_get_indexdef(idx.indexrelid)
  231. END AS exprdef,
  232. CASE am.amname
  233. WHEN %s THEN
  234. CASE (option & 1)
  235. WHEN 1 THEN 'DESC' ELSE 'ASC'
  236. END
  237. END as ordering,
  238. c2.reloptions as attoptions
  239. FROM (
  240. SELECT *
  241. FROM
  242. pg_index i,
  243. unnest(i.indkey, i.indoption)
  244. WITH ORDINALITY koi(key, option, arridx)
  245. ) idx
  246. LEFT JOIN pg_class c ON idx.indrelid = c.oid
  247. LEFT JOIN pg_class c2 ON idx.indexrelid = c2.oid
  248. LEFT JOIN pg_am am ON c2.relam = am.oid
  249. LEFT JOIN
  250. pg_attribute attr ON attr.attrelid = c.oid AND attr.attnum = idx.key
  251. WHERE c.relname = %s AND pg_catalog.pg_table_is_visible(c.oid)
  252. ) s2
  253. GROUP BY indexname, indisunique, indisprimary, amname, exprdef, attoptions;
  254. """,
  255. [self.index_default_access_method, table_name],
  256. )
  257. for (
  258. index,
  259. columns,
  260. unique,
  261. primary,
  262. orders,
  263. type_,
  264. definition,
  265. options,
  266. ) in cursor.fetchall():
  267. if index not in constraints:
  268. basic_index = (
  269. type_ == self.index_default_access_method
  270. and
  271. # '_btree' references
  272. # django.contrib.postgres.indexes.BTreeIndex.suffix.
  273. not index.endswith("_btree")
  274. and options is None
  275. )
  276. constraints[index] = {
  277. "columns": columns if columns != [None] else [],
  278. "orders": orders if orders != [None] else [],
  279. "primary_key": primary,
  280. "unique": unique,
  281. "foreign_key": None,
  282. "check": False,
  283. "index": True,
  284. "type": Index.suffix if basic_index else type_,
  285. "definition": definition,
  286. "options": options,
  287. }
  288. return constraints