check database in auth_query for multitenant deployments

.github/workflows/submit.yml

@@ -4,7 +4,7 @@ on:
     branches:
       - main
 env:
-  VERSION: "0.2.3"
+  VERSION: "0.2.4"
   IMAGE_NAME: "lanterndata/lantern-suite"
 jobs:
   docker:

Dockerfile

@@ -8,7 +8,17 @@ ARG TARGETARCH
 ARG PG_CRON_VERSION="7e91e72b1bebc5869bb900d9253cc9e92518b33f"
 ENV OS_ARCH="${TARGETARCH:-amd64}"
 
-RUN apt update && apt install -y curl wget make jq pgbouncer procps bc git-all gcc postgresql-server-dev-${PG_VERSION}
+RUN apt update && apt install -y autoconf automake libtool pandoc libevent-dev pkg-config curl wget make jq procps bc git-all gcc postgresql-server-dev-${PG_VERSION}
+
+# Install pgbouncer
+RUN git clone https://github.com/var77/pgbouncer.git /tmp/pgbouncer && \
+    cd /tmp/pgbouncer && \
+    git submodule init && git submodule update && \
+    ./autogen.sh && \
+    ./configure --prefix=/usr/local && \
+    make -j && \
+    make install && \
+    rm -rf /tmp/pgbouncer
 
 # Install pg_cron
 RUN git clone https://github.com/citusdata/pg_cron.git /tmp/pg_cron && \
@@ -23,7 +33,7 @@ RUN git clone --branch v0.7.3-lanterncloud https://github.com/lanterndata/pgvect
     make OPTFLAGS="" -j && \
     make install
 
-    # Install Lantern
+# Install Lantern
 RUN cd /tmp && \
     wget https://github.com/lanterndata/lantern/releases/download/v${LANTERN_VERSION}/lantern-${LANTERN_VERSION}.tar -O lantern.tar && \
     tar xf lantern.tar && \

scripts/configure-pgbouncer.sh

@@ -11,6 +11,24 @@ cd $PG_SOCKET_DIR
 check_cmd="pg_isready"
 ready_counter=$POSTGRESQL_INIT_MAX_TIMEOUT
 
+########################
+# Check if the provided argument is a boolean or is the string 'yes/true'
+# Arguments:
+#   $1 - Value to check
+# Returns:
+#   Boolean
+#########################
+is_boolean_yes() {
+    local -r bool="${1:-}"
+    # comparison is performed without regard to the case of alphabetic characters
+    shopt -s nocasematch
+    if [[ "$bool" = 1 || "$bool" =~ ^(yes|true)$ ]]; then
+        true
+    else
+        false
+    fi
+}
+
 generate_random_string() {
     local count="32"
     local filter
@@ -52,6 +70,18 @@ SQL
 
 psql -Atq -U postgres -d postgres -c "SELECT concat('\"', usename, '\" \"', passwd, '\"') FROM pg_shadow WHERE usename='_pgbouncer'" > userlist.txt
 
+auth_query='SELECT usename, passwd FROM pg_shadow WHERE usename=$1'
+
+if is_boolean_yes "${LANTERN_MULTI_TENANT:-no}"; then
+  psql -U postgres -h $PG_SOCKET_DIR -p $POSTGRESQL_PORT_NUMBER postgres -t <<SQL 
+    BEGIN;
+      CREATE TABLE IF NOT EXISTS _pgbouncer_auth(id SERIAL PRIMARY KEY, rolname NAME, dbname NAME);
+      GRANT SELECT ON _pgbouncer_auth TO _pgbouncer;
+    COMMIT;
+SQL
+auth_query='SELECT psh.usename, psh.passwd FROM pg_shadow psh INNER JOIN pg_user pgu ON pgu.usename=psh.usename LEFT JOIN _pgbouncer_auth auth ON auth.rolname=psh.usename WHERE psh.usename=$1 AND (pgu.usesuper=TRUE OR auth.dbname=$2);'
+fi
+
 cat <<EOF > pgbouncer.ini
 [pgbouncer]
 # Connection settings
@@ -62,7 +92,7 @@ auth_user = _pgbouncer
 auth_hba_file = $POSTGRESQL_PGHBA_FILE
 auth_file = userlist.txt
 unix_socket_dir = $PG_SOCKET_DIR
-auth_query = SELECT usename, passwd FROM pg_shadow WHERE usename=\$1
+auth_query = $auth_query
 auth_dbname = postgres
 
 pidfile = pgbouncer.pid