[wip] working through changes for EnsureUploaded to support postgresql.

Author: simukka

dbintf.go

@@ -3,6 +3,7 @@ package sqlcode
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 )
 
 type DB interface {
@@ -11,6 +12,7 @@ type DB interface {
 	QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row
 	Conn(ctx context.Context) (*sql.Conn, error)
 	BeginTx(ctx context.Context, txOptions *sql.TxOptions) (*sql.Tx, error)
+	Driver() driver.Driver
 }
 
 var _ DB = &sql.DB{}

dbops.go

@@ -3,11 +3,25 @@ package sqlcode
 import (
 	"context"
 	"database/sql"
+
+	mssql "github.com/denisenkom/go-mssqldb"
+	"github.com/jackc/pgx/v5/stdlib"
 )
 
 func Exists(ctx context.Context, dbc DB, schemasuffix string) (bool, error) {
 	var schemaID int
-	err := dbc.QueryRowContext(ctx, `select isnull(schema_id(@p1), 0)`, SchemaName(schemasuffix)).Scan(&schemaID)
+
+	driver := dbc.Driver()
+	var qs string
+
+	if _, ok := driver.(*mssql.Driver); ok {
+		qs = `select isnull(schema_id(@p1), 0)`
+	}
+	if _, ok := driver.(*stdlib.Driver); ok {
+		qs = `select coalesce((select oid from pg_namespace where nspname = $1),0)`
+	}
+
+	err := dbc.QueryRowContext(ctx, qs, SchemaName(schemasuffix)).Scan(&schemaID)
 	if err != nil {
 		return false, err
 	}

deployable.go

@@ -11,6 +11,9 @@ import (
 	"time"
 
 	mssql "github.com/denisenkom/go-mssqldb"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/stdlib"
+	pgxstdlib "github.com/jackc/pgx/v5/stdlib"
 	"github.com/vippsas/sqlcode/sqlparser"
 )
 
@@ -77,21 +80,22 @@ func impersonate(ctx context.Context, dbc DB, username string, f func(conn *sql.
 // Upload will create and upload the schema; resulting in an error
 // if the schema already exists
 func (d *Deployable) Upload(ctx context.Context, dbc DB) error {
-	// First, impersonate a user with minimal privileges to get at least
-	// some level of sandboxing so that migration scripts can't do anything
-	// the caller didn't expect them to.
-	return impersonate(ctx, dbc, "sqlcode-deploy-sandbox-user", func(conn *sql.Conn) error {
+	driver := dbc.Driver()
+	qs := make(map[string][]interface{}, 1)
+
+	var uploadFunc = func(conn *sql.Conn) error {
 		tx, err := conn.BeginTx(ctx, nil)
 		if err != nil {
 			return err
 		}
 
-		_, err = tx.ExecContext(ctx, `sqlcode.CreateCodeSchema`,
-			sql.Named("schemasuffix", d.SchemaSuffix),
-		)
-		if err != nil {
-			_ = tx.Rollback()
-			return err
+		for q, args := range qs {
+			_, err = tx.ExecContext(ctx, q, args...)
+
+			if err != nil {
+				_ = tx.Rollback()
+				return fmt.Errorf("failed to execute (%s) with arg(%s) in schema %s: %w", q, args, d.SchemaSuffix, err)
+			}
 		}
 
 		preprocessed, err := Preprocess(d.CodeBase, d.SchemaSuffix)
@@ -123,8 +127,36 @@ func (d *Deployable) Upload(ctx context.Context, dbc DB) error {
 
 		return nil
 
-	})
+	}
+
+	if _, ok := driver.(*mssql.Driver); ok {
+		// First, impersonate a user with minimal privileges to get at least
+		// some level of sandboxing so that migration scripts can't do anything
+		// the caller didn't expect them to.
+		qs["sqlcode.CreateCodeSchema"] = []interface {
+		}{
+			sql.Named("schemasuffix", d.SchemaSuffix),
+		}
 
+		return impersonate(ctx, dbc, "sqlcode-deploy-sandbox-user", uploadFunc)
+	}
+
+	if _, ok := driver.(*stdlib.Driver); ok {
+		qs[`set role "sqlcode-deploy-sandbox-user"`] = nil
+		qs[`call sqlcode.createcodeschema(@schemasuffix)`] = []interface{}{
+			pgx.NamedArgs{"schemasuffix": d.SchemaSuffix},
+		}
+		conn, err := dbc.Conn(ctx)
+		if err != nil {
+			return err
+		}
+		defer func() {
+			_ = conn.Close()
+		}()
+		return uploadFunc(conn)
+	}
+
+	return fmt.Errorf("failed to determine sql driver to upload schema: %s", d.SchemaSuffix)
 }
 
 // EnsureUploaded checks that the schema with the suffix already exists,
@@ -137,37 +169,51 @@ func (d *Deployable) EnsureUploaded(ctx context.Context, dbc DB) error {
 		return nil
 	}
 
+	driver := dbc.Driver()
 	lockResourceName := "sqlcode.EnsureUploaded/" + d.SchemaSuffix
 
+	var lockRetCode int
+	var lockQs string
+	var unlockQs string
+	var err error
+
 	// When a lock is opened with the Transaction lock owner,
 	// that lock is released when the transaction is committed or rolled back.
-	var lockRetCode int
-	err := dbc.QueryRowContext(ctx, `
-declare @retcode int;
-exec @retcode = sp_getapplock @Resource = @resource, @LockMode = 'Shared', @LockOwner = 'Session', @LockTimeout = @timeoutMs;
-select @retcode;
-`,
-		sql.Named("resource", lockResourceName),
-		sql.Named("timeoutMs", 20000),
-	).Scan(&lockRetCode)
+	if _, ok := driver.(*pgxstdlib.Driver); ok {
+		lockQs = `select sqlcode.get_applock(@resource, @timeout)`
+		unlockQs = `select sqlcode.release_applock(@resource)`
+
+		err = dbc.QueryRowContext(ctx, lockQs, pgx.NamedArgs{
+			"resource":  lockResourceName,
+			"timeoutMs": 20000,
+		}).Scan(&lockRetCode)
+
+		defer func() {
+			dbc.ExecContext(ctx, unlockQs, pgx.NamedArgs{"resource": lockResourceName})
+		}()
+	}
+
+	if _, ok := driver.(*mssql.Driver); ok {
+		// TODO
+
+		defer func() {
+			// TODO: This returns an error if the lock is already released
+			_, _ = dbc.ExecContext(ctx, unlockQs,
+				sql.Named("Resource", lockResourceName),
+				sql.Named("LockOwner", "Session"),
+			)
+		}()
+	}
+
 	if err != nil {
 		return err
 	}
 	if lockRetCode < 0 {
 		return errors.New("was not able to get lock before timeout")
 	}
-
-	defer func() {
-		// TODO: This returns an error if the lock is already released
-		_, _ = dbc.ExecContext(ctx, `sp_releaseapplock`,
-			sql.Named("Resource", lockResourceName),
-			sql.Named("LockOwner", "Session"),
-		)
-	}()
-
 	exists, err := Exists(ctx, dbc, d.SchemaSuffix)
 	if err != nil {
-		return err
+		return fmt.Errorf("unable to determine if schema %s exists: %w", d.SchemaSuffix, err)
 	}
 
 	if exists {

docker-compose.pgsql.yml

@@ -7,6 +7,7 @@ services:
       POSTGRES_PASSWORD: VippsPw1
       POSTGRES_USER: sa
       POSTGRES_DB: master
+      PGOPTIONS: "-c log_error_verbosity=verbose -c log_statement=all"
     healthcheck:
       test: ["CMD-SHELL", "pg_isready", "-d", "db_prod"]
       interval: 1s

migrations/0003.sqlcode.pgsql

@@ -156,6 +156,52 @@ exception
 end;
 $$;
 
+-- similar behaviour as mssql getapplock
+-- PostgreSQL advisory locks are session-based by default
+create or replace function sqlcode.get_applock(
+    resource text,
+    timeout_ms integer default 0
+)
+returns integer
+language plpgsql
+as $$
+declare
+    resource_key bigint;
+    acquired     boolean;
+    waited_ms    integer := 0;
+begin
+    -- convert string to advisory-lock key
+    select hashtext(resource) into resource_key;
+
+    -- attempt lock with timeout loop
+    loop
+        select pg_try_advisory_lock_shared(resource_key)
+        into acquired;
+
+        if acquired then
+            return 1;  -- lock acquired (success)
+        end if;
+
+        if waited_ms >= timeout_ms then
+            return 0;  -- timeout
+        end if;
+
+        perform pg_sleep(0.01);    -- sleep 10 ms
+        waited_ms := waited_ms + 10;
+    end loop;
+
+    return null;  -- safety fallback (should never hit)
+end;
+$$;
+
+create or replace function sqlcode.release_applock(resource text)
+returns boolean
+language sql
+as $$
+    select pg_advisory_unlock_shared(hashtext(resource));
+$$;
+
+
 -- ensure procedures are owned by the definer role
 alter procedure sqlcode.createcodeschema(varchar)
     owner to "sqlcode-definer-role";

sqltest/sqlcode_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -58,5 +59,25 @@ func Test_EnsureUploaded(t *testing.T) {
 		}
 
 		fixture.RunMigrationFile("../migrations/0003.sqlcode.pgsql")
+
+		ctx := context.Background()
+
+		_, err := fixture.adminDB.Exec(`grant create on database @database to "sqlcode-definer-role"`,
+			pgx.NamedArgs{"database": fixture.DBName})
+		require.NoError(t, err)
+
+		require.NoError(t, SQL.EnsureUploaded(ctx, fixture.DB))
+		patched := SQL.Patch(`[code].Test`)
+
+		res, err := fixture.DB.ExecContext(ctx, patched)
+		require.NoError(t, err)
+		rowsAffected, err := res.RowsAffected()
+		require.NoError(t, err)
+		assert.Equal(t, int64(1), rowsAffected)
+
+		schemas := SQL.ListUploaded(ctx, fixture.DB)
+		require.Len(t, schemas, 1)
+		require.Equal(t, 6, schemas[0].Objects)
+		require.Equal(t, "5420c0269aaf", schemas[0].Suffix())
 	})
 }