1#!/usr/local/bin/python3.8
2# -*- coding: utf-8 -*-
3
4# Copyright: (c) 2019, Andrew Klychkov (@Andersson007) <aaklychkov@mail.ru>
5# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
6
7from __future__ import absolute_import, division, print_function
8__metaclass__ = type
9
10ANSIBLE_METADATA = {
11    'metadata_version': '1.1',
12    'status': ['preview'],
13    'supported_by': 'community'
14}
15
16DOCUMENTATION = r'''
17---
18module: postgresql_table
19short_description: Create, drop, or modify a PostgreSQL table
20description:
21- Allows to create, drop, rename, truncate a table, or change some table attributes.
22version_added: '2.8'
23options:
24  table:
25    description:
26    - Table name.
27    required: true
28    aliases:
29    - name
30    type: str
31  state:
32    description:
33    - The table state. I(state=absent) is mutually exclusive with I(tablespace), I(owner), I(unlogged),
34      I(like), I(including), I(columns), I(truncate), I(storage_params) and, I(rename).
35    type: str
36    default: present
37    choices: [ absent, present ]
38  tablespace:
39    description:
40    - Set a tablespace for the table.
41    required: false
42    type: str
43  owner:
44    description:
45    - Set a table owner.
46    type: str
47  unlogged:
48    description:
49    - Create an unlogged table.
50    type: bool
51    default: no
52  like:
53    description:
54    - Create a table like another table (with similar DDL).
55      Mutually exclusive with I(columns), I(rename), and I(truncate).
56    type: str
57  including:
58    description:
59    - Keywords that are used with like parameter, may be DEFAULTS, CONSTRAINTS, INDEXES, STORAGE, COMMENTS or ALL.
60      Needs I(like) specified. Mutually exclusive with I(columns), I(rename), and I(truncate).
61    type: str
62  columns:
63    description:
64    - Columns that are needed.
65    type: list
66    elements: str
67  rename:
68    description:
69    - New table name. Mutually exclusive with I(tablespace), I(owner),
70      I(unlogged), I(like), I(including), I(columns), I(truncate), and I(storage_params).
71    type: str
72  truncate:
73    description:
74    - Truncate a table. Mutually exclusive with I(tablespace), I(owner), I(unlogged),
75      I(like), I(including), I(columns), I(rename), and I(storage_params).
76    type: bool
77    default: no
78  storage_params:
79    description:
80    - Storage parameters like fillfactor, autovacuum_vacuum_treshold, etc.
81      Mutually exclusive with I(rename) and I(truncate).
82    type: list
83    elements: str
84  db:
85    description:
86    - Name of database to connect and where the table will be created.
87    type: str
88    aliases:
89    - login_db
90  session_role:
91    description:
92    - Switch to session_role after connecting.
93      The specified session_role must be a role that the current login_user is a member of.
94    - Permissions checking for SQL commands is carried out as though
95      the session_role were the one that had logged in originally.
96    type: str
97  cascade:
98    description:
99    - Automatically drop objects that depend on the table (such as views).
100      Used with I(state=absent) only.
101    type: bool
102    default: no
103    version_added: '2.9'
104notes:
105- If you do not pass db parameter, tables will be created in the database
106  named postgres.
107- PostgreSQL allows to create columnless table, so columns param is optional.
108- Unlogged tables are available from PostgreSQL server version 9.1.
109seealso:
110- module: postgresql_sequence
111- module: postgresql_idx
112- module: postgresql_info
113- module: postgresql_tablespace
114- module: postgresql_owner
115- module: postgresql_privs
116- module: postgresql_copy
117- name: CREATE TABLE reference
118  description: Complete reference of the CREATE TABLE command documentation.
119  link: https://www.postgresql.org/docs/current/sql-createtable.html
120- name: ALTER TABLE reference
121  description: Complete reference of the ALTER TABLE  command documentation.
122  link: https://www.postgresql.org/docs/current/sql-altertable.html
123- name: DROP TABLE reference
124  description: Complete reference of the DROP TABLE command documentation.
125  link: https://www.postgresql.org/docs/current/sql-droptable.html
126- name: PostgreSQL data types
127  description: Complete reference of the PostgreSQL data types documentation.
128  link: https://www.postgresql.org/docs/current/datatype.html
129author:
130- Andrei Klychkov (@Andersson007)
131extends_documentation_fragment: postgres
132'''
133
134EXAMPLES = r'''
135- name: Create tbl2 in the acme database with the DDL like tbl1 with testuser as an owner
136  postgresql_table:
137    db: acme
138    name: tbl2
139    like: tbl1
140    owner: testuser
141
142- name: Create tbl2 in the acme database and tablespace ssd with the DDL like tbl1 including comments and indexes
143  postgresql_table:
144    db: acme
145    table: tbl2
146    like: tbl1
147    including: comments, indexes
148    tablespace: ssd
149
150- name: Create test_table with several columns in ssd tablespace with fillfactor=10 and autovacuum_analyze_threshold=1
151  postgresql_table:
152    name: test_table
153    columns:
154    - id bigserial primary key
155    - num bigint
156    - stories text
157    tablespace: ssd
158    storage_params:
159    - fillfactor=10
160    - autovacuum_analyze_threshold=1
161
162- name: Create an unlogged table in schema acme
163  postgresql_table:
164    name: acme.useless_data
165    columns: waste_id int
166    unlogged: true
167
168- name: Rename table foo to bar
169  postgresql_table:
170    table: foo
171    rename: bar
172
173- name: Rename table foo from schema acme to bar
174  postgresql_table:
175    name: acme.foo
176    rename: bar
177
178- name: Set owner to someuser
179  postgresql_table:
180    name: foo
181    owner: someuser
182
183- name: Change tablespace of foo table to new_tablespace and set owner to new_user
184  postgresql_table:
185    name: foo
186    tablespace: new_tablespace
187    owner: new_user
188
189- name: Truncate table foo
190  postgresql_table:
191    name: foo
192    truncate: yes
193
194- name: Drop table foo from schema acme
195  postgresql_table:
196    name: acme.foo
197    state: absent
198
199- name: Drop table bar cascade
200  postgresql_table:
201    name: bar
202    state: absent
203    cascade: yes
204'''
205
206RETURN = r'''
207table:
208  description: Name of a table.
209  returned: always
210  type: str
211  sample: 'foo'
212state:
213  description: Table state.
214  returned: always
215  type: str
216  sample: 'present'
217owner:
218  description: Table owner.
219  returned: always
220  type: str
221  sample: 'postgres'
222tablespace:
223  description: Tablespace.
224  returned: always
225  type: str
226  sample: 'ssd_tablespace'
227queries:
228  description: List of executed queries.
229  returned: always
230  type: str
231  sample: [ 'CREATE TABLE "test_table" (id bigint)' ]
232storage_params:
233  description: Storage parameters.
234  returned: always
235  type: list
236  sample: [ "fillfactor=100", "autovacuum_analyze_threshold=1" ]
237'''
238
239try:
240    from psycopg2.extras import DictCursor
241except ImportError:
242    # psycopg2 is checked by connect_to_db()
243    # from ansible.module_utils.postgres
244    pass
245
246from ansible.module_utils.basic import AnsibleModule
247from ansible.module_utils.database import pg_quote_identifier
248from ansible.module_utils.postgres import (
249    connect_to_db,
250    exec_sql,
251    get_conn_params,
252    postgres_common_argument_spec,
253)
254
255
256# ===========================================
257# PostgreSQL module specific support methods.
258#
259
260class Table(object):
261    def __init__(self, name, module, cursor):
262        self.name = name
263        self.module = module
264        self.cursor = cursor
265        self.info = {
266            'owner': '',
267            'tblspace': '',
268            'storage_params': [],
269        }
270        self.exists = False
271        self.__exists_in_db()
272        self.executed_queries = []
273
274    def get_info(self):
275        """Getter to refresh and get table info"""
276        self.__exists_in_db()
277
278    def __exists_in_db(self):
279        """Check table exists and refresh info"""
280        if "." in self.name:
281            schema = self.name.split('.')[-2]
282            tblname = self.name.split('.')[-1]
283        else:
284            schema = 'public'
285            tblname = self.name
286
287        query = ("SELECT t.tableowner, t.tablespace, c.reloptions "
288                 "FROM pg_tables AS t "
289                 "INNER JOIN pg_class AS c ON  c.relname = t.tablename "
290                 "INNER JOIN pg_namespace AS n ON c.relnamespace = n.oid "
291                 "WHERE t.tablename = %(tblname)s "
292                 "AND n.nspname = %(schema)s")
293        res = exec_sql(self, query, query_params={'tblname': tblname, 'schema': schema},
294                       add_to_executed=False)
295        if res:
296            self.exists = True
297            self.info = dict(
298                owner=res[0][0],
299                tblspace=res[0][1] if res[0][1] else '',
300                storage_params=res[0][2] if res[0][2] else [],
301            )
302
303            return True
304        else:
305            self.exists = False
306            return False
307
308    def create(self, columns='', params='', tblspace='',
309               unlogged=False, owner=''):
310        """
311        Create table.
312        If table exists, check passed args (params, tblspace, owner) and,
313        if they're different from current, change them.
314        Arguments:
315        params - storage params (passed by "WITH (...)" in SQL),
316            comma separated.
317        tblspace - tablespace.
318        owner - table owner.
319        unlogged - create unlogged table.
320        columns - column string (comma separated).
321        """
322        name = pg_quote_identifier(self.name, 'table')
323
324        changed = False
325
326        if self.exists:
327            if tblspace == 'pg_default' and self.info['tblspace'] is None:
328                pass  # Because they have the same meaning
329            elif tblspace and self.info['tblspace'] != tblspace:
330                self.set_tblspace(tblspace)
331                changed = True
332
333            if owner and self.info['owner'] != owner:
334                self.set_owner(owner)
335                changed = True
336
337            if params:
338                param_list = [p.strip(' ') for p in params.split(',')]
339
340                new_param = False
341                for p in param_list:
342                    if p not in self.info['storage_params']:
343                        new_param = True
344
345                if new_param:
346                    self.set_stor_params(params)
347                    changed = True
348
349            if changed:
350                return True
351            return False
352
353        query = "CREATE"
354        if unlogged:
355            query += " UNLOGGED TABLE %s" % name
356        else:
357            query += " TABLE %s" % name
358
359        if columns:
360            query += " (%s)" % columns
361        else:
362            query += " ()"
363
364        if params:
365            query += " WITH (%s)" % params
366
367        if tblspace:
368            query += " TABLESPACE %s" % pg_quote_identifier(tblspace, 'database')
369
370        if exec_sql(self, query, ddl=True):
371            changed = True
372
373        if owner:
374            changed = self.set_owner(owner)
375
376        return changed
377
378    def create_like(self, src_table, including='', tblspace='',
379                    unlogged=False, params='', owner=''):
380        """
381        Create table like another table (with similar DDL).
382        Arguments:
383        src_table - source table.
384        including - corresponds to optional INCLUDING expression
385            in CREATE TABLE ... LIKE statement.
386        params - storage params (passed by "WITH (...)" in SQL),
387            comma separated.
388        tblspace - tablespace.
389        owner - table owner.
390        unlogged - create unlogged table.
391        """
392        changed = False
393
394        name = pg_quote_identifier(self.name, 'table')
395
396        query = "CREATE"
397        if unlogged:
398            query += " UNLOGGED TABLE %s" % name
399        else:
400            query += " TABLE %s" % name
401
402        query += " (LIKE %s" % pg_quote_identifier(src_table, 'table')
403
404        if including:
405            including = including.split(',')
406            for i in including:
407                query += " INCLUDING %s" % i
408
409        query += ')'
410
411        if params:
412            query += " WITH (%s)" % params
413
414        if tblspace:
415            query += " TABLESPACE %s" % pg_quote_identifier(tblspace, 'database')
416
417        if exec_sql(self, query, ddl=True):
418            changed = True
419
420        if owner:
421            changed = self.set_owner(owner)
422
423        return changed
424
425    def truncate(self):
426        query = "TRUNCATE TABLE %s" % pg_quote_identifier(self.name, 'table')
427        return exec_sql(self, query, ddl=True)
428
429    def rename(self, newname):
430        query = "ALTER TABLE %s RENAME TO %s" % (pg_quote_identifier(self.name, 'table'),
431                                                 pg_quote_identifier(newname, 'table'))
432        return exec_sql(self, query, ddl=True)
433
434    def set_owner(self, username):
435        query = "ALTER TABLE %s OWNER TO %s" % (pg_quote_identifier(self.name, 'table'),
436                                                pg_quote_identifier(username, 'role'))
437        return exec_sql(self, query, ddl=True)
438
439    def drop(self, cascade=False):
440        if not self.exists:
441            return False
442
443        query = "DROP TABLE %s" % pg_quote_identifier(self.name, 'table')
444        if cascade:
445            query += " CASCADE"
446        return exec_sql(self, query, ddl=True)
447
448    def set_tblspace(self, tblspace):
449        query = "ALTER TABLE %s SET TABLESPACE %s" % (pg_quote_identifier(self.name, 'table'),
450                                                      pg_quote_identifier(tblspace, 'database'))
451        return exec_sql(self, query, ddl=True)
452
453    def set_stor_params(self, params):
454        query = "ALTER TABLE %s SET (%s)" % (pg_quote_identifier(self.name, 'table'), params)
455        return exec_sql(self, query, ddl=True)
456
457
458# ===========================================
459# Module execution.
460#
461
462
463def main():
464    argument_spec = postgres_common_argument_spec()
465    argument_spec.update(
466        table=dict(type='str', required=True, aliases=['name']),
467        state=dict(type='str', default="present", choices=["absent", "present"]),
468        db=dict(type='str', default='', aliases=['login_db']),
469        tablespace=dict(type='str'),
470        owner=dict(type='str'),
471        unlogged=dict(type='bool', default=False),
472        like=dict(type='str'),
473        including=dict(type='str'),
474        rename=dict(type='str'),
475        truncate=dict(type='bool', default=False),
476        columns=dict(type='list', elements='str'),
477        storage_params=dict(type='list', elements='str'),
478        session_role=dict(type='str'),
479        cascade=dict(type='bool', default=False),
480    )
481    module = AnsibleModule(
482        argument_spec=argument_spec,
483        supports_check_mode=True,
484    )
485
486    table = module.params["table"]
487    state = module.params["state"]
488    tablespace = module.params["tablespace"]
489    owner = module.params["owner"]
490    unlogged = module.params["unlogged"]
491    like = module.params["like"]
492    including = module.params["including"]
493    newname = module.params["rename"]
494    storage_params = module.params["storage_params"]
495    truncate = module.params["truncate"]
496    columns = module.params["columns"]
497    cascade = module.params["cascade"]
498
499    if state == 'present' and cascade:
500        module.warn("cascade=true is ignored when state=present")
501
502    # Check mutual exclusive parameters:
503    if state == 'absent' and (truncate or newname or columns or tablespace or like or storage_params or unlogged or owner or including):
504        module.fail_json(msg="%s: state=absent is mutually exclusive with: "
505                             "truncate, rename, columns, tablespace, "
506                             "including, like, storage_params, unlogged, owner" % table)
507
508    if truncate and (newname or columns or like or unlogged or storage_params or owner or tablespace or including):
509        module.fail_json(msg="%s: truncate is mutually exclusive with: "
510                             "rename, columns, like, unlogged, including, "
511                             "storage_params, owner, tablespace" % table)
512
513    if newname and (columns or like or unlogged or storage_params or owner or tablespace or including):
514        module.fail_json(msg="%s: rename is mutually exclusive with: "
515                             "columns, like, unlogged, including, "
516                             "storage_params, owner, tablespace" % table)
517
518    if like and columns:
519        module.fail_json(msg="%s: like and columns params are mutually exclusive" % table)
520    if including and not like:
521        module.fail_json(msg="%s: including param needs like param specified" % table)
522
523    conn_params = get_conn_params(module, module.params)
524    db_connection = connect_to_db(module, conn_params, autocommit=False)
525    cursor = db_connection.cursor(cursor_factory=DictCursor)
526
527    if storage_params:
528        storage_params = ','.join(storage_params)
529
530    if columns:
531        columns = ','.join(columns)
532
533    ##############
534    # Do main job:
535    table_obj = Table(table, module, cursor)
536
537    # Set default returned values:
538    changed = False
539    kw = {}
540    kw['table'] = table
541    kw['state'] = ''
542    if table_obj.exists:
543        kw = dict(
544            table=table,
545            state='present',
546            owner=table_obj.info['owner'],
547            tablespace=table_obj.info['tblspace'],
548            storage_params=table_obj.info['storage_params'],
549        )
550
551    if state == 'absent':
552        changed = table_obj.drop(cascade=cascade)
553
554    elif truncate:
555        changed = table_obj.truncate()
556
557    elif newname:
558        changed = table_obj.rename(newname)
559        q = table_obj.executed_queries
560        table_obj = Table(newname, module, cursor)
561        table_obj.executed_queries = q
562
563    elif state == 'present' and not like:
564        changed = table_obj.create(columns, storage_params,
565                                   tablespace, unlogged, owner)
566
567    elif state == 'present' and like:
568        changed = table_obj.create_like(like, including, tablespace,
569                                        unlogged, storage_params)
570
571    if changed:
572        if module.check_mode:
573            db_connection.rollback()
574        else:
575            db_connection.commit()
576
577        # Refresh table info for RETURN.
578        # Note, if table has been renamed, it gets info by newname:
579        table_obj.get_info()
580        db_connection.commit()
581        if table_obj.exists:
582            kw = dict(
583                table=table,
584                state='present',
585                owner=table_obj.info['owner'],
586                tablespace=table_obj.info['tblspace'],
587                storage_params=table_obj.info['storage_params'],
588            )
589        else:
590            # We just change the table state here
591            # to keep other information about the dropped table:
592            kw['state'] = 'absent'
593
594    kw['queries'] = table_obj.executed_queries
595    kw['changed'] = changed
596    db_connection.close()
597    module.exit_json(**kw)
598
599
600if __name__ == '__main__':
601    main()
602