diff --git a/pkg/format/pgre/pgre.go b/pkg/format/pgre/pgre.go index 44334ee..78a20dc 100644 --- a/pkg/format/pgre/pgre.go +++ b/pkg/format/pgre/pgre.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "net/url" + "strings" "github.com/go-pg/pg/v10" "github.com/vmkteam/pgdesigner/pkg/pgd" @@ -50,7 +51,7 @@ func connectDB(dsn string) (*pg.DB, error) { // IsDSN returns true if the input looks like a PostgreSQL DSN. func IsDSN(s string) bool { - return len(s) > 11 && (s[:11] == "postgres://" || s[:13] == "postgresql://") + return strings.HasPrefix(s, "postgres://") || strings.HasPrefix(s, "postgresql://") } func parseDSN(dsn string) (*pg.Options, error) { diff --git a/pkg/format/pgre/pgre_test.go b/pkg/format/pgre/pgre_test.go index 4565ba3..dfd947a 100644 --- a/pkg/format/pgre/pgre_test.go +++ b/pkg/format/pgre/pgre_test.go @@ -81,6 +81,25 @@ func skipIfNoPG(t *testing.T) { } } +func TestIsDSN(t *testing.T) { + tests := []struct { + name string + in string + want bool + }{ + {name: "postgres scheme", in: "postgres://localhost/db", want: true}, + {name: "postgresql scheme", in: "postgresql://localhost/db", want: true}, + {name: "plain path", in: "docs/apisrv.xml", want: false}, + {name: "trailing dot 12 chars", in: "docs/apisrv.", want: false}, + {name: "empty", in: "", want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, IsDSN(tt.in)) + }) + } +} + func TestParseIndexColumns(t *testing.T) { tests := []struct { name string