1 # -*- coding: utf-8 -*-

     2 #

     3 # PgBrowser - browse database schema and metadata

     4 #

     5 # Some of the queries came from psql.

     6 #

     7 # Copyright (c) 2011  Radek Brich <radek.brich@devl.cz>

     8 #

     9 # Permission is hereby granted, free of charge, to any person obtaining a copy

    10 # of this software and associated documentation files (the "Software"), to deal

    11 # in the Software without restriction, including without limitation the rights

    12 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

    13 # copies of the Software, and to permit persons to whom the Software is

    14 # furnished to do so, subject to the following conditions:

    15 #

    16 # The above copyright notice and this permission notice shall be included in

    17 # all copies or substantial portions of the Software.

    18 #

    19 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

    20 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

    21 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

    22 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

    23 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

    24 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

    25 # THE SOFTWARE.

    26 

    27 

    28 from collections import OrderedDict

    29 

    30 

    31 class Column:

    32     def __init__(self, browser, table,

    33             name, type, notnull, hasdefault, default, description):

    34         self.browser = browser  # Browser instance

    35         self.table = table  # Table instance

    36         self.name = name

    37         self.type = type

    38         self.notnull = notnull

    39         self.hasdefault = hasdefault

    40         self.default = default

    41         self.description = description

    42 

    43 

    44 class Constraint:

    45     def __init__(self, browser, table, name, type, fname, fschema, definition):

    46         self.browser = browser

    47         self.table = table

    48         self.name = name

    49         self.type = type

    50         self.fname = fname  # foreign table name

    51         self.fschema = fschema  # foreign table schema

    52         self.definition = definition

    53 

    54 

    55 class Index:

    56     def __init__(self, browser, table,

    57             name, primary, unique, clustered, valid, definition,

    58             columns, size):

    59         self.browser = browser

    60         self.table = table

    61         self.name = name

    62         self.primary = primary

    63         self.unique = unique

    64         self.clustered = clustered

    65         self.valid = valid

    66         self.definition = definition

    67         self.columns = columns

    68         self.size = size

    69 

    70 

    71 class Table:

    72     def __init__(self, browser, schema, name, owner, size, description, options):

    73         self._columns = None

    74         self._constraints = None

    75         self._indexes = None

    76         self.browser = browser  # Browser instance

    77         self.schema = schema  # Schema instance

    78         self.name = name  # table name, str

    79         self.owner = owner

    80         self.size = size

    81         self.description = description

    82         self.options = options or []

    83 

    84     def refresh(self):

    85         self.refresh_columns()

    86         self.refresh_constraints()

    87         self.refresh_indexes()

    88 

    89     def refresh_columns(self):

    90         rows = self.browser.list_columns(self.name, self.schema.name)

    91         self._columns = OrderedDict([(x['name'], Column(self.browser, self, **x)) for x in rows])

    92 

    93     def refresh_constraints(self):

    94         rows = self.browser.list_constraints(self.name, self.schema.name)

    95         self._constraints = OrderedDict([(x['name'], Constraint(self.browser, self, **x)) for x in rows])

    96 

    97     def refresh_indexes(self):

    98         rows = self.browser.list_indexes(self.name, self.schema.name)

    99         self._indexes = OrderedDict([(x['name'], Index(self.browser, self, **x)) for x in rows])

   100 

   101     def getcolumns(self):

   102         if self._columns is None:

   103             self.refresh_columns()

   104         return self._columns

   105     columns = property(getcolumns)

   106 

   107     def getconstraints(self):

   108         if self._constraints is None:

   109             self.refresh_constraints()

   110         return self._constraints

   111     constraints = property(getconstraints)

   112 

   113     def getindexes(self):

   114         if self._indexes is None:

   115             self.refresh_indexes()

   116         return self._indexes

   117     indexes = property(getindexes)

   118 

   119 

   120 class Argument:

   121     def __init__(self, browser, function, name, type, mode, default):

   122         # PgBrowser instance

   123         self.browser = browser

   124         # Function instance

   125         self.function = function

   126         self.name = name

   127         self.type = type

   128         self.mode = mode

   129         self.default = default

   130 

   131 

   132 class Function:

   133     def __init__(self, browser, schema, oid, name, function_name, type, result, source):

   134         self.browser = browser

   135         self.schema = schema

   136         self.oid = oid

   137         #: unique name - function name + arg types

   138         self.name = name

   139         #: pure function name without args

   140         self.function_name = function_name

   141         self.type = type

   142         self.result = result

   143         self.source = source

   144         self._arguments = None

   145         self._definition = None

   146 

   147     def refresh(self):

   148         self.refresh_args()

   149 

   150     def refresh_args(self):

   151         rows = self.browser.list_function_args(self.oid)

   152         self._arguments = OrderedDict([(x['name'], Argument(self.browser, self, **x)) for x in rows])

   153 

   154     @property

   155     def arguments(self):

   156         if self._arguments is None:

   157             self.refresh_args()

   158         return self._arguments

   159 

   160     @property

   161     def definition(self):

   162         """Get full function definition including CREATE command."""

   163         if not self._definition:

   164             self._definition = self.browser.get_function_definition(self.oid)

   165         return self._definition

   166 

   167 

   168 class Type:

   169     def __init__(self, browser, schema, name, type, elements, description):

   170         self.browser = browser

   171         self.schema = schema

   172         self.name = name

   173         self.type = type

   174         self.elements = elements

   175         self.description = description

   176 

   177 

   178 class Schema:

   179     def __init__(self, browser, name, owner, acl, description, system):

   180         self._tables = None

   181         self._functions = None

   182         self._types = None

   183         self.browser = browser

   184         self.name = name

   185         self.owner = owner

   186         self.acl = acl

   187         self.description = description

   188         self.system = system

   189 

   190     def refresh(self):

   191         self.refresh_tables()

   192         self.refresh_functions()

   193 

   194     def refresh_tables(self):

   195         rows = self.browser.list_tables(self.name)

   196         self._tables = OrderedDict([(x['name'], Table(self.browser, self, **x)) for x in rows])

   197 

   198     def refresh_functions(self):

   199         rows = self.browser.list_functions(self.name)

   200         self._functions = OrderedDict([(x['name'], Function(self.browser, self, **x)) for x in rows])

   201 

   202     def refresh_types(self):

   203         rows = self.browser.list_types(self.name)

   204         self._types = OrderedDict([(x['name'], Type(self.browser, self, **x)) for x in rows])

   205 

   206     @property

   207     def tables(self):

   208         if self._tables is None:

   209             self.refresh_tables()

   210         return self._tables

   211 

   212     @property

   213     def functions(self):

   214         if self._functions is None:

   215             self.refresh_functions()

   216         return self._functions

   217 

   218     @property

   219     def types(self):

   220         if self._types is None:

   221             self.refresh_types()

   222         return self._types

   223 

   224 

   225 class PgBrowser:

   226     def __init__(self, conn=None):

   227         self._schemas = None

   228         self.conn = conn

   229 

   230     def setconn(self, conn=None):

   231         self.conn = conn

   232 

   233     def refresh(self):

   234         self.refresh_schemas()

   235 

   236     def refresh_schemas(self):

   237         rows = self.list_schemas()

   238         self._schemas = OrderedDict([(x['name'], Schema(self, **x)) for x in rows])

   239 

   240     @property

   241     def schemas(self):

   242         if self._schemas is None:

   243             self.refresh_schemas()

   244         return self._schemas

   245 

   246     def _query(self, query, args):

   247         try:

   248             curs = self.conn.cursor()

   249             curs.execute(query, args)

   250             curs.connection.commit()

   251             rows = curs.fetchall()

   252             return [dict(zip([desc[0] for desc in curs.description], row)) for row in rows]

   253         finally:

   254             curs.close()

   255 

   256     def list_databases(self):

   257         return self._query('''

   258             SELECT

   259                 d.datname as "name",

   260                 pg_catalog.pg_get_userbyid(d.datdba) as "owner",

   261                 pg_catalog.pg_encoding_to_char(d.encoding) as "encoding",

   262                 d.datcollate as "collation",

   263                 d.datctype as "ctype",

   264                 d.datacl AS "acl",

   265                 CASE WHEN pg_catalog.has_database_privilege(d.datname, 'CONNECT')

   266                     THEN pg_catalog.pg_database_size(d.datname)

   267                     ELSE -1 -- No access

   268                 END as "size",

   269                 t.spcname as "tablespace",

   270                 pg_catalog.shobj_description(d.oid, 'pg_database') as "description"

   271             FROM pg_catalog.pg_database d

   272             JOIN pg_catalog.pg_tablespace t on d.dattablespace = t.oid

   273             ORDER BY 1;

   274             ''', [])

   275 

   276     def list_schemas(self):

   277         return self._query('''

   278             SELECT

   279                 n.nspname AS "name",

   280                 pg_catalog.pg_get_userbyid(n.nspowner) AS "owner",

   281                 n.nspacl AS "acl",

   282                 pg_catalog.obj_description(n.oid, 'pg_namespace') AS "description",

   283                 CASE WHEN n.nspname IN ('information_schema', 'pg_catalog', 'pg_toast')

   284                     OR n.nspname ~ '^pg_temp_' OR n.nspname ~ '^pg_toast_temp_'

   285                     THEN TRUE

   286                     ELSE FALSE

   287                 END AS "system"

   288             FROM pg_catalog.pg_namespace n

   289             ORDER BY 1;

   290             ''', [])

   291 

   292     def list_tables(self, schema='public'):

   293         return self._query('''

   294             SELECT

   295                 c.relname as "name",

   296                 pg_catalog.pg_get_userbyid(c.relowner) as "owner",

   297                 pg_catalog.pg_relation_size(c.oid) as "size",

   298                 pg_catalog.obj_description(c.oid, 'pg_class') as "description",

   299                 c.reloptions as "options"

   300             FROM pg_catalog.pg_class c

   301             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace

   302             WHERE n.nspname = %s AND c.relkind IN ('r','s','')

   303             ORDER BY 1;

   304             ''', [schema])

   305 

   306     def list_columns(self, table, schema='public', order=2):

   307         return self._query('''

   308             SELECT

   309                 --a.attrelid,

   310                 a.attname as "name",

   311                 format_type(a.atttypid, a.atttypmod) AS "type",

   312                 a.attnotnull as "notnull",

   313                 a.atthasdef as "hasdefault",

   314                 pg_catalog.pg_get_expr(d.adbin, d.adrelid) as "default",

   315                 pg_catalog.col_description(a.attrelid, a.attnum) AS "description"

   316             FROM pg_catalog.pg_attribute a

   317             LEFT JOIN pg_catalog.pg_class c ON a.attrelid = c.oid

   318             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace

   319             LEFT JOIN pg_catalog.pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum

   320             WHERE n.nspname = %s AND c.relname = %s AND a.attnum > 0 AND NOT a.attisdropped

   321             ORDER BY ''' + str(order), [schema, table])

   322 

   323     def list_constraints(self, table, schema='public'):

   324         return self._query('''

   325             SELECT

   326                 r.conname AS "name",

   327                 r.contype AS "type",

   328                 cf.relname AS "fname",

   329                 nf.nspname AS "fschema",

   330                 pg_catalog.pg_get_constraintdef(r.oid, true) as "definition"

   331             FROM pg_catalog.pg_constraint r

   332             JOIN pg_catalog.pg_class c ON r.conrelid = c.oid

   333             JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace

   334             LEFT JOIN pg_catalog.pg_class cf ON r.confrelid = cf.oid

   335             LEFT JOIN pg_catalog.pg_namespace nf ON nf.oid = cf.relnamespace

   336             WHERE n.nspname = %s AND c.relname = %s

   337             ORDER BY 1

   338             ''', [schema, table])

   339 

   340     def list_indexes(self, table, schema='public'):

   341         return self._query('''

   342             SELECT

   343                 c2.relname as "name",

   344                 i.indisprimary as "primary",

   345                 i.indisunique as "unique",

   346                 i.indisclustered as "clustered",

   347                 i.indisvalid as "valid",

   348                 pg_catalog.pg_get_indexdef(i.indexrelid, 0, true) as "definition",

   349                 ARRAY(SELECT a.attname FROM pg_catalog.pg_attribute a WHERE a.attrelid = c2.oid ORDER BY attnum) AS "columns",

   350                 pg_catalog.pg_relation_size(c2.oid) as "size"

   351                 --c2.reltablespace as "tablespace_oid"

   352             FROM pg_catalog.pg_class c

   353             JOIN pg_catalog.pg_index i ON c.oid = i.indrelid

   354             JOIN pg_catalog.pg_class c2 ON i.indexrelid = c2.oid

   355             JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace

   356             WHERE n.nspname = %(schema)s AND c.relname = %(table)s

   357             ORDER BY i.indisprimary DESC, i.indisunique DESC, c2.relname

   358             ''', {'schema': schema, 'table': table})

   359 

   360     def list_functions(self, schema='public'):

   361         '''List functions in schema.'''

   362         return self._query('''

   363             SELECT

   364                 p.oid as "oid",

   365                 p.proname || '(' || array_to_string(

   366                     array(SELECT pg_catalog.format_type(unnest(p.proargtypes), NULL)),

   367                       ', '

   368                 ) || ')' as "name",

   369                 p.proname as "function_name",

   370                 pg_catalog.pg_get_function_result(p.oid) as "result",

   371                 p.prosrc as "source",

   372                 CASE

   373                     WHEN p.proisagg THEN 'agg'

   374                     WHEN p.proiswindow THEN 'window'

   375                     WHEN p.prorettype = 'pg_catalog.trigger'::pg_catalog.regtype THEN 'trigger'

   376                     ELSE 'normal'

   377                 END as "type"

   378             FROM pg_catalog.pg_proc p

   379             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace

   380             WHERE n.nspname = %s

   381             ORDER BY 1, 2, 4;

   382             ''', [schema])

   383 

   384     def list_function_args(self, oid):

   385         """List function arguments.

   386 

   387         Notes about query:

   388             type: Use allargtypes if present, argtypes otherwise.

   389                 The trick with [0:999] moves lower bound from 0 to default 1

   390                 by slicing all elements (slices has always lower bound 1).

   391             mode: This trick makes array of NULLs of same length as argnames,

   392                 in case argmodes is NULL.

   393             default: Use pg_get_expr, split output by ', '

   394                 FIXME: will fail if ', ' is present in default value string.

   395         """

   396         return self._query('''

   397             SELECT

   398               unnest(p.proargnames) AS "name",

   399               pg_catalog.format_type(unnest(

   400                 COALESCE(p.proallargtypes, (p.proargtypes::oid[])[0:999])

   401               ), NULL) AS "type",

   402               unnest(

   403                 COALESCE(

   404                   p.proargmodes::text[],

   405                   array(SELECT NULL::text FROM generate_series(1, array_upper(p.proargnames, 1)))

   406                 )

   407               ) AS "mode",

   408               unnest(array_cat(

   409                 array_fill(NULL::text, array[COALESCE(array_upper(p.proargnames,1),0) - p.pronargdefaults]),

   410                 string_to_array(pg_get_expr(p.proargdefaults, 'pg_proc'::regclass, true), ', ')

   411               )) AS "default"

   412             FROM pg_proc p

   413             WHERE p.oid = %s''', [oid])

   414 

   415     def get_function_definition(self, oid):

   416         """Get full function definition, including CREATE command etc.

   417 

   418         Args:

   419             oid: function oid from pg_catalog.pg_proc (returned by list_functions)

   420 

   421         """

   422         return self._query('''SELECT pg_get_functiondef(%s) AS definition;''', [oid])[0]['definition']

   423 

   424     def list_types(self, schema='public'):

   425         """List types in schema.

   426 

   427         http://www.postgresql.org/docs/8.4/static/catalog-pg-type.html

   428 

   429         """

   430         return self._query('''

   431             SELECT

   432                 t.typname AS "name",

   433                 CASE

   434                     WHEN t.typtype = 'b' THEN 'base'::text

   435                     WHEN t.typtype = 'c' THEN 'composite'::text

   436                     WHEN t.typtype = 'd' THEN 'domain'::text

   437                     WHEN t.typtype = 'e' THEN 'enum'::text

   438                     WHEN t.typtype = 'p' THEN 'pseudo'::text

   439                 END AS "type",

   440                 ARRAY(

   441                       SELECT e.enumlabel

   442                       FROM pg_catalog.pg_enum e

   443                       WHERE e.enumtypid = t.oid

   444                       ORDER BY e.oid

   445                 ) AS "elements",

   446                 pg_catalog.obj_description(t.oid, 'pg_type') AS "description"

   447             FROM pg_catalog.pg_type t

   448             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace

   449             WHERE (t.typrelid = 0 OR (SELECT c.relkind = 'c' FROM pg_catalog.pg_class c WHERE c.oid = t.typrelid))

   450               AND NOT EXISTS(SELECT 1 FROM pg_catalog.pg_type el WHERE el.oid = t.typelem AND el.typarray = t.oid)

   451                   AND n.nspname <> 'pg_catalog'

   452                   AND n.nspname <> 'information_schema'

   453               AND n.nspname = %(schema)s

   454             ORDER BY 1, 2;

   455         ''', {'schema': schema})

   456 

   457     def list_sequences(self, schema=None):

   458         '''List sequences in schema.'''

   459         return self._query('''

   460             SELECT

   461                 nc.nspname AS "sequence_schema",

   462                 c.relname AS "sequence_name",

   463                 t.relname AS "related_table",

   464                 a.attname AS "related_column",

   465                 format_type(a.atttypid, a.atttypmod) AS "related_column_type"

   466             FROM pg_class c

   467             JOIN pg_namespace nc ON nc.oid = c.relnamespace

   468             JOIN pg_depend d ON d.objid = c.oid

   469             JOIN pg_class t ON d.refobjid = t.oid

   470             JOIN pg_attribute a ON (d.refobjid, d.refobjsubid) = (a.attrelid, a.attnum)

   471             WHERE c.relkind = 'S' AND NOT pg_is_other_temp_schema(nc.oid)

   472             ''' + (schema and ' AND nc.nspname = %(schema)s' or '') + '''

   473         ''', {'schema': schema})

   474 

   475     def list_column_usage(self, table, column, schema='public'):

   476         '''List objects using the column.

   477 

   478         Currently shows views and constraints which use the column.

   479 

   480         This is useful to find which views block alteration of column type etc.

   481 

   482         '''

   483         return self._query('''

   484             SELECT

   485                 'view' AS type, view_schema AS schema, view_name AS name

   486             FROM information_schema.view_column_usage

   487             WHERE table_schema=%(schema)s AND table_name=%(table)s AND column_name=%(column)s

   488 

   489             UNION

   490 

   491             SELECT

   492                 'constraint' AS type, constraint_schema AS schema, constraint_name AS name

   493             FROM information_schema.constraint_column_usage

   494             WHERE table_schema=%(schema)s AND table_name=%(table)s AND column_name=%(column)s

   495             ''', {'schema':schema, 'table':table, 'column':column})

   496