1import os
2
3import infra.basetest
4
5
6class TestPostgreSQL(infra.basetest.BRTest):
7    # We use a specific configuration for:
8    # - using Aarch64, to have more than 256MB memory,
9    # - to have an ext4 rootfs image exposed as a virtio storage
10    #   (rather than cpio initrd). This will save some memory, as the
11    #   rootfs image is big.
12    config = \
13        """
14        BR2_aarch64=y
15        BR2_TOOLCHAIN_EXTERNAL=y
16        BR2_TARGET_GENERIC_GETTY_PORT="ttyAMA0"
17        BR2_LINUX_KERNEL=y
18        BR2_LINUX_KERNEL_CUSTOM_VERSION=y
19        BR2_LINUX_KERNEL_CUSTOM_VERSION_VALUE="6.1.72"
20        BR2_LINUX_KERNEL_USE_CUSTOM_CONFIG=y
21        BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE="board/qemu/aarch64-virt/linux.config"
22        BR2_PACKAGE_POSTGRESQL=y
23        BR2_TARGET_ROOTFS_EXT2=y
24        BR2_TARGET_ROOTFS_EXT2_4=y
25        BR2_TARGET_ROOTFS_EXT2_SIZE="256M"
26        # BR2_TARGET_ROOTFS_TAR is not set
27        """
28
29    def __init__(self, names):
30        super(TestPostgreSQL, self).__init__(names)
31        self.db_admin = "postgres"
32        self.db_user = "br_user"
33        self.db_name = "br_database"
34        self.backup_file = "dump.sql"
35        self.pgdata_dir = "/var/lib/pgsql"
36
37    def run_user_db_query(self, user, database, query, opts=None):
38        cmd = f"psql --username={user} --dbname={database}"
39        cmd += f' --command="{query}"'
40        if opts is not None:
41            cmd += " " + opts
42        self.assertRunOk(cmd)
43
44    def run_admin_sql_query(self, query, opts=None):
45        self.run_user_db_query(self.db_admin, self.db_admin, query, opts)
46
47    def run_sql_query(self, query, opts=None):
48        self.run_user_db_query(self.db_user, self.db_name, query, opts)
49
50    def cleanup_database(self):
51        # This cleanup is useful when run-test -k is used. It makes
52        # this test idempotent. Since the drive storage is preserved
53        # between reboots, this cleanup will prevent errors during the
54        # user/db creation.
55
56        # Drop the test database, if it exists.
57        cmd = f"dropdb --username={self.db_admin} "
58        cmd += f"--if-exists {self.db_name}"
59        self.assertRunOk(cmd)
60
61        # Drop the test user, if it exists.
62        cmd = f"dropuser --username={self.db_admin} "
63        cmd += f"--if-exists {self.db_user}"
64        self.assertRunOk(cmd)
65
66    def create_tables(self):
67        sql_query = "CREATE TABLE fruits ("
68        sql_query += "id integer PRIMARY KEY, name varchar(16) NOT NULL);"
69        self.run_sql_query(sql_query)
70
71        sql_query = "CREATE TABLE colors ("
72        sql_query += "id integer PRIMARY KEY, name varchar(16) NOT NULL);"
73        self.run_sql_query(sql_query)
74
75        sql_query = "CREATE TABLE fruit_colors ("
76        sql_query += "fruit_id integer REFERENCES fruits(id), "
77        sql_query += "color_id integer REFERENCES colors(id), "
78        sql_query += "UNIQUE (fruit_id, color_id));"
79        self.run_sql_query(sql_query)
80
81    def insert_data(self):
82        fruits = ["Banana", "Blueberry", "Orange", "Raspberry"]
83        fruit_id = 1
84        for fruit in fruits:
85            sql_query = "INSERT INTO fruits (id, name) "
86            sql_query += f"VALUES ({fruit_id}, '{fruit}');"
87            self.run_sql_query(sql_query)
88            fruit_id += 1
89
90        colors = ["Blue", "Orange", "Red", "Yellow"]
91        color_id = 1
92        for color in colors:
93            sql_query = "INSERT INTO colors (id, name) "
94            sql_query += f"VALUES ({color_id}, '{color}');"
95            self.run_sql_query(sql_query)
96            color_id += 1
97
98        fruit_colors = [(1, 4), (2, 1), (3, 2), (4, 3)]
99        for fruit_color in fruit_colors:
100            fruit_id, color_id = fruit_color
101            sql_query = "INSERT INTO fruit_colors (fruit_id, color_id) "
102            sql_query += f"VALUES ({fruit_id}, {color_id});"
103            self.run_sql_query(sql_query)
104
105    def query_database(self):
106        sql_query = "SELECT "
107        sql_query += "fruits.name AS fruit, colors.name AS color "
108        sql_query += "FROM fruits, colors, fruit_colors "
109        sql_query += "WHERE fruits.id = fruit_colors.fruit_id "
110        sql_query += "AND colors.id = fruit_colors.color_id "
111        sql_query += "ORDER BY fruit;"
112        self.run_sql_query(sql_query)
113
114    def test_run(self):
115        drive = os.path.join(self.builddir, "images", "rootfs.ext4")
116        kern = os.path.join(self.builddir, "images", "Image")
117        self.emulator.boot(arch="aarch64",
118                           kernel=kern,
119                           kernel_cmdline=["root=/dev/vda console=ttyAMA0"],
120                           options=["-M", "virt",
121                                    "-cpu", "cortex-a57",
122                                    "-m", "512M",
123                                    "-smp", "2",
124                                    "-drive", f"file={drive},if=virtio,format=raw"])
125        self.emulator.login()
126
127        # Check the server binary can execute.
128        self.assertRunOk("postgres --version")
129
130        # Check the client binary can execute.
131        self.assertRunOk("psql --version")
132
133        # Check the server is ready.
134        self.assertRunOk("pg_isready")
135
136        # Query the server version from the client.
137        sql_query = "SELECT version();"
138        self.run_admin_sql_query(sql_query,
139                                 opts="--tuples-only --no-align")
140
141        self.cleanup_database()
142
143        # Create a new user.
144        cmd = f"createuser --username={self.db_admin} "
145        cmd += "--no-superuser --no-createdb --no-createrole "
146        cmd += self.db_user
147        self.assertRunOk(cmd)
148
149        # Create a new database.
150        cmd = f"createdb --username={self.db_admin} "
151        cmd += f"--owner={self.db_user} "
152        cmd += f'{self.db_name} "Test Database for Buildroot Test"'
153        self.assertRunOk(cmd)
154
155        self.create_tables()
156
157        self.insert_data()
158
159        self.query_database()
160
161        # Update a table.
162        sql_query = "UPDATE fruits SET name = 'Lemon' WHERE id = 1;"
163        self.run_sql_query(sql_query)
164
165        # Backup the test database.
166        cmd = f"pg_dump --username={self.db_user} --dbname={self.db_name} "
167        cmd += f"--file={self.backup_file} --inserts"
168        self.assertRunOk(cmd)
169
170        # Drop all the tables.
171        sql_query = "DROP TABLE fruit_colors, fruits, colors;"
172        self.run_sql_query(sql_query)
173
174        # Query the server status.
175        cmd = f"su - {self.db_admin} -c 'pg_ctl status -D {self.pgdata_dir}'"
176        self.assertRunOk(cmd)
177
178        # Stop the server.
179        cmd = f"su - {self.db_admin} -c 'pg_ctl stop -D {self.pgdata_dir}'"
180        self.assertRunOk(cmd)
181
182        # Check the server is no longer ready.
183        _, exit_code = self.emulator.run("pg_isready")
184        self.assertNotEqual(exit_code, 0)
185
186        # Restart the server.
187        cmd = f"su - {self.db_admin} -c 'pg_ctl start -D {self.pgdata_dir}'"
188        self.assertRunOk(cmd)
189
190        # Restore the backup.
191        cmd = f"psql --username={self.db_user} --dbname={self.db_name} "
192        cmd += f"--file={self.backup_file}"
193        self.assertRunOk(cmd)
194
195        # Query on last time our data, to check the backup restoration
196        # succeeded.
197        self.query_database()
198