PostgreSQL: Install PL/Python, add basic function

This installs PL/Python (specifically plpython3u), enables it in the
database, and creates a function id36_to_id that calls the Python
function with the same name inside the tildes.lib.id module. This will
enable doing queries similar to this, when I have a topic's ID36 from
the site:

SELECT * FROM topics WHERE topic_id = id36_to_id('asdf');

The fact that this was possible to set up without having to port the
id36_to_id logic to a different language is blowing my mind a little.
There are some really interesting possibilities from being able to
import all of the Python code into the database itself.

salt/salt/final-setup.sls

@@ -4,6 +4,7 @@ initialize-db:
   cmd.run:
     - name: {{ bin_dir }}/python -c "from scripts.initialize_db import initialize_db; initialize_db('{{ app_dir }}/{{ pillar['ini_file'] }}')"
     - cwd: {{ app_dir }}
+    - runas: postgres
     - require:
       - postgres_database: tildes
     - unless: psql -U tildes tildes -c "SELECT user_id FROM users;"

salt/salt/postgresql/init.sls

@@ -1,3 +1,5 @@
+{% from 'common.jinja2' import app_dir, venv_dir %}
+
 postgresql:
   pkgrepo.managed:
     - name: deb http://apt.postgresql.org/pub/repos/apt/ xenial-pgdg main
@@ -16,6 +18,13 @@ postgresql:
     - watch:
       - file: /etc/postgresql/{{ pillar['postgresql_version'] }}/main/*.conf
 
+install-plpython3:
+  pkg.installed:
+    - name: postgresql-plpython3-{{ pillar['postgresql_version'] }}
+    - refresh: True
+    - require:
+      - pkgrepo: postgresql
+
 set-lock-timeout:
   file.replace:
     - name: /etc/postgresql/{{ pillar['postgresql_version'] }}/main/postgresql.conf
@@ -60,3 +69,17 @@ enable-pg-stat-statements:
     - mode: 640
   require:
     - service: postgresql
+
+# set PYTHONPATH env var in postgresql so PL/Python can access all the modules
+set-postgresql-pythonpath:
+  file.managed:
+    - name: /etc/postgresql/{{ pillar['postgresql_version'] }}/main/environment
+    - contents: "PYTHONPATH='{{ venv_dir }}/lib/python3.7/site-packages:{{ app_dir }}'"
+    - user: postgres
+    - group: postgres
+    - mode: 644
+    - watch_in:
+      - module: set-postgresql-pythonpath
+  module.wait:
+    - service.restart:
+      - name: postgresql.service

salt/salt/postgresql/site-db.sls

@@ -1,6 +1,8 @@
 site-db-user:
   postgres_user.present:
     - name: tildes
+    # set the tildes user to a database superuser in the dev environment only
+    - superuser: {{ grains['id'] == 'dev' }}
     - require:
       - service: postgresql
 
@@ -46,6 +48,13 @@ site-db-enable-pg_trgm:
     - require:
       - postgres_database: tildes
 
+site-db-enable-plpython3u:
+  postgres_extension.present:
+    - name: plpython3u
+    - maintenance_db: tildes
+    - require:
+      - postgres_database: tildes
+
 site-db-pg_hba-lines:
   file.accumulated:
     - name: pg_hba_lines

salt/salt/postgresql/test-db.sls

@@ -33,6 +33,13 @@ test-db-enable-pg_trgm:
     - require:
       - postgres_database: tildes_test
 
+test-db-enable-plpython3u:
+  postgres_extension.present:
+    - name: plpython3u
+    - maintenance_db: tildes_test
+    - require:
+      - postgres_database: tildes_test
+
 test-db-pg_hba-lines:
   file.accumulated:
     - name: pg_hba_lines

tildes/scripts/initialize_db.py

@@ -21,7 +21,12 @@ from tildes.models.user import User
 
 
 def initialize_db(config_path: str, alembic_config_path: Optional[str] = None) -> None:
-    """Load the app config and create the database tables."""
+    """Load the app config and create the database tables.
+
+    This function will probably only complete successfully when run as user postgres,
+    since the run_sql_scripts_in_dir method runs psql through that user to be able to
+    create functions using untrusted languages (PL/Python specifically).
+    """
     db_session = get_session_from_config(config_path)
     engine = db_session.bind
 
@@ -57,7 +62,7 @@ def create_tables(connectable: Connectable) -> None:
 
 
 def run_sql_scripts_in_dir(path: str, engine: Engine) -> None:
-    """Run all sql scripts in a directory."""
+    """Run all sql scripts in a directory, as the database superuser."""
     for root, _, files in os.walk(path):
         sql_files = [filename for filename in files if filename.endswith(".sql")]
         for sql_file in sql_files:
@@ -65,7 +70,7 @@ def run_sql_scripts_in_dir(path: str, engine: Engine) -> None:
                 [
                     "psql",
                     "-U",
-                    engine.url.username,
+                    "postgres",
                     "-f",
                     os.path.join(root, sql_file),
                     engine.url.database,

tildes/sql/init/functions/utils.sql

@@ -0,0 +1,8 @@
+-- Copyright (c) 2019 Tildes contributors <code@tildes.net>
+-- SPDX-License-Identifier: AGPL-3.0-or-later
+
+CREATE OR REPLACE FUNCTION id36_to_id(id36 TEXT) RETURNS INTEGER AS $$
+    from tildes.lib.id import id36_to_id
+
+    return id36_to_id(id36)
+$$ IMMUTABLE LANGUAGE plpython3u;