Repository: danvergara/dblab (go)

Issue:
[FEATURE] Multi schemas
**Is your feature request related to a problem? Please describe.**
I have a database with multiple schemas, but I am unable to select the schemas that I need.

**Describe the solution you'd like**
I think it could be great to have the option to select the schemas that we want to show of the database, or maybe could be shown in a panel with the schemas.

**Describe alternatives you've considered**
Currently, I think there is no workaround to work with multiple schemas

**Additional context**
This is an example with the following schemas:
<img width="1619" alt="image" src="https://user-images.githubusercontent.com/20668624/179146608-dc2aae59-32b3-4286-b026-4db0d2c3d4fe.png">
<img width="189" alt="image" src="https://user-images.githubusercontent.com/20668624/179146694-76865132-b02e-41a3-ad63-f2d529db9f9e.png">
I think currently we are only showing the public schema:
<img width="341" alt="image" src="https://user-images.githubusercontent.com/20668624/179146791-d98c75fc-a7e2-44e2-b35c-80df21f2a9b8.png">
I was looking for examples and I found this:
<img width="569" alt="image" src="https://user-images.githubusercontent.com/20668624/179147898-2840c46f-9ec0-4b7a-9f7c-24efc353077f.png">
We can select the schemas that we need and are shown like this:
<img width="295" alt="image" src="https://user-images.githubusercontent.com/20668624/179148018-feb5ef20-cabd-4b06-8c7c-1f5c4c763441.png">



Pull request description:
Multi schema
## Description

@israteneda pointed out that we need some sort of support for dealing with multiple schemas, since Postgres has this enabled.

Fixes #95

## Type of change

Please delete options that are not relevant.

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [x] This change requires a documentation update

## How Has This Been Tested?

Updated the tests to support the schema feature. I test it out on my local system, as well.

## Checklist:

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have checked my code and corrected any misspellings


Patch:
diff --git a/.dblab.yaml b/.dblab.yaml
index d9e592e..cf09a83 100644
--- a/.dblab.yaml
+++ b/.dblab.yaml
@@ -4,5 +4,6 @@ database:
   db: "users"
   password: "password"
   user: "postgres"
+  schema: "public"
   driver: "postgres"
 limit: 50
diff --git a/Makefile b/Makefile
index 6c9b659..173e64c 100644
--- a/Makefile
+++ b/Makefile
@@ -44,7 +44,7 @@ build-sqlite3:
 .PHONY: run
 ## run: Runs the application
 run: build
-	./dblab --host localhost --user postgres --db users --pass password --ssl disable --port 5432 --driver postgres --limit 50
+	./dblab --host localhost --user postgres --db users --pass password --schema public --ssl disable --port 5432 --driver postgres --limit 50
 
 .PHONY: run-mysql
 ## run-mysql: Runs the application with a connection to mysql
@@ -53,7 +53,7 @@ run-mysql: build
 
 .PHONY: run-sqlite3
 ## run-sqlite3: Runs the application with a connection to sqlite3
-run-sqlite3: build
+run-sqlite3: build-sqlite3
 	docker-compose run --rm dblab-sqlite3
 	./dblab --db db/dblab.db --driver sqlite3
 
@@ -77,6 +77,11 @@ down:
 form: build
 	./dblab
 
+.PHONY: create-migration
+## create: Creates golang-migrate migration files
+create-migration:
+	 migrate create -ext sql -dir ./db/migrations $(file_name)
+
 .PHONY: help
 ## help: Prints this help message
 help:
diff --git a/README.md b/README.md
index c9678e7..35fc810 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,7 @@ Flags:
       --limit int       Size of the result set from the table content query (default 100)
       --pass string     Password for user
       --port string     Server port
+      --schema string   Database schema (postgres only)
       --ssl string      SSL mode
   -u, --url string      Database connection string
       --user string     Database user
@@ -94,6 +95,7 @@ Use "dblab [command] --help" for more information about a command.
 You can start the app passing no flags or parameters, you'll be asked for connection data instead.
 ![Alt Text](screenshots/dblab-default-form.gif)
 
+
 ```sh
 $ dblab --host localhost --user myuser --db users --pass password --ssl disable --port 5432 --driver postgres --limit 50
 $ dblab --db path/to/file.sqlite3 --driver sqlite3
@@ -107,6 +109,13 @@ $ dblab --url mysql://user:password@tcp(host:port)/db
 $ dblab --url file:test.db?cache=shared&mode=memory
 ```
 
+if you're using PostgreSQL, you have the option to define the schema you want to work with, the default value is `public`.
+
+```sh
+$ dblab --host localhost --user myuser --db users --pass password --schema myschema --ssl disable --port 5432 --driver postgres --limit 50
+$ dblab --url postgres://user:password@host:port/database?sslmode=[mode] --schema myschema
+```
+
 Now, you can use a configuration file to make a connection to the database.
 
 ```sh
@@ -125,6 +134,10 @@ database:
   password: "password"
   user: "postgres"
   driver: "postgres"
+  # optional
+  # postgres only
+  # default value: public
+  schema: "myschema"
 limit: 50
 ```
 
diff --git a/cmd/root.go b/cmd/root.go
index 466d04f..93887e6 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -19,6 +19,7 @@ var (
 	port   string
 	user   string
 	pass   string
+	schema string
 	db     string
 	ssl    string
 	limit  int
@@ -48,6 +49,7 @@ func NewRootCmd() *cobra.Command {
 					User:   user,
 					Pass:   pass,
 					DBName: db,
+					Schema: schema,
 					SSL:    ssl,
 					Limit:  limit,
 				}
@@ -107,6 +109,7 @@ func init() {
 	rootCmd.Flags().StringVarP(&user, "user", "", "", "Database user")
 	rootCmd.Flags().StringVarP(&pass, "pass", "", "", "Password for user")
 	rootCmd.Flags().StringVarP(&db, "db", "", "", "Database name")
+	rootCmd.Flags().StringVarP(&schema, "schema", "", "", "Database schema (postgres only)")
 	rootCmd.Flags().StringVarP(&ssl, "ssl", "", "", "SSL mode")
 	rootCmd.Flags().IntVarP(&limit, "limit", "", 100, "Size of the result set from the table content query")
 }
diff --git a/db/migrations/20221011024901_create_myschema_user_table.down.sql b/db/migrations/20221011024901_create_myschema_user_table.down.sql
new file mode 100644
index 0000000..c80488b
--- /dev/null
+++ b/db/migrations/20221011024901_create_myschema_user_table.down.sql
@@ -0,0 +1,1 @@
+DROP TABLE myschema.users;
diff --git a/db/migrations/20221011024901_create_myschema_user_table.up.sql b/db/migrations/20221011024901_create_myschema_user_table.up.sql
new file mode 100644
index 0000000..552beba
--- /dev/null
+++ b/db/migrations/20221011024901_create_myschema_user_table.up.sql
@@ -0,0 +1,5 @@
+CREATE TABLE myschema.users
+(
+    id SERIAL PRIMARY KEY,
+    username VARCHAR(100) NOT NULL UNIQUE
+);
diff --git a/docker-compose.yml b/docker-compose.yml
index 6091ffd..e606630 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -39,6 +39,7 @@ services:
       - DB_NAME=users
       - DB_PORT=5432
       - DB_DRIVER=postgres
+      - DB_SCHEMA=public
     entrypoint: ["/bin/bash", "./scripts/entrypoint.dev.sh"]
     networks:
       - dblab
diff --git a/pkg/client/client.go b/pkg/client/client.go
index b74e9a9..27a4f21 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -21,7 +21,7 @@ import (
 // Client is used to store the pool of db connection.
 type Client struct {
 	db                *sqlx.DB
-	driver            string
+	driver, schema    string
 	paginationManager *pagination.Manager
 	limit             int
 }
@@ -44,6 +44,21 @@ func New(opts command.Options) (*Client, error) {
 		limit:  opts.Limit,
 	}
 
+	if opts.Schema == "" {
+		c.schema = "public"
+	} else {
+		c.schema = opts.Schema
+	}
+
+	switch c.driver {
+	case "postgres":
+		fallthrough
+	case "postgresql":
+		if _, err = db.Exec(fmt.Sprintf("set search_path='%s'", c.schema)); err != nil {
+			return nil, err
+		}
+	}
+
 	pm, err := pagination.New(c.limit, 0, "")
 	if err != nil {
 		return nil, err
@@ -185,22 +200,28 @@ func (c *Client) TotalPages() int {
 
 // ShowTables list all the tables in the database on the tables panel.
 func (c *Client) ShowTables() ([]string, error) {
-	var query string
+	var (
+		query string
+		err   error
+		args  []interface{}
+	)
+
 	tables := make([]string, 0)
 
 	switch c.driver {
 	case "postgres":
 		fallthrough
 	case "postgresql":
-		query = `
-		SELECT
-			table_name
-		FROM
-			information_schema.tables
-		WHERE
-			table_schema = 'public'
-		ORDER BY
-			table_name;`
+		psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
+		query, args, err = psql.Select("table_name").
+			From("information_schema.tables").
+			Where(sq.Eq{"table_schema": c.schema}).
+			OrderBy("table_name").
+			ToSql()
+		if err != nil {
+			return nil, err
+		}
+
 	case "mysql":
 		query = "SHOW TABLES;"
 	case "sqlite3":
@@ -214,7 +235,7 @@ func (c *Client) ShowTables() ([]string, error) {
 			name NOT LIKE 'sqlite_%';`
 	}
 
-	rows, err := c.db.Queryx(query)
+	rows, err := c.db.Queryx(query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -301,7 +322,7 @@ func (c *Client) tableContent(tableName string) ([][]string, []string, error) {
 	var query string
 
 	if c.driver == "postgres" || c.driver == "postgresql" {
-		query = fmt.Sprintf("SELECT * FROM public.%s LIMIT %d OFFSET %d;", tableName, c.paginationManager.Limit(), c.paginationManager.Offset())
+		query = fmt.Sprintf("SELECT * FROM %q LIMIT %d OFFSET %d;", tableName, c.paginationManager.Limit(), c.paginationManager.Offset())
 	} else {
 		query = fmt.Sprintf("SELECT * FROM %s LIMIT %d OFFSET %d;", tableName, c.paginationManager.Limit(), c.paginationManager.Offset())
 	}
@@ -317,7 +338,7 @@ func (c *Client) tableCount(tableName string) (int, error) {
 	)
 
 	if c.driver == "postgres" || c.driver == "postgresql" {
-		query = fmt.Sprintf("SELECT COUNT(*) FROM public.%s;", tableName)
+		query = fmt.Sprintf("SELECT COUNT(*) FROM %q;", tableName)
 	} else {
 		query = fmt.Sprintf("SELECT COUNT(*) FROM %s;", tableName)
 	}
@@ -337,33 +358,42 @@ func (c *Client) tableStructure(tableName string) ([][]string, []string, error)
 	case "postgres":
 		fallthrough
 	case "postgresql":
-		query = `
-        SELECT
-			c.column_name,
-			c.is_nullable,
-			c.data_type,
-			c.character_maximum_length,
-			c.numeric_precision,
-			c.numeric_scale,
-			c.ordinal_position,
-			tc.constraint_type pkey
-		FROM
-			information_schema.columns c
-		LEFT JOIN
-			information_schema.constraint_column_usage AS ccu
-		ON
-			c.table_schema = ccu.table_schema
-			AND c.table_name = ccu.table_name
-			AND c.column_name = ccu.column_name
-		LEFT JOIN
-			information_schema.table_constraints AS tc
-		ON
-			ccu.constraint_schema = tc.constraint_schema
-			and ccu.constraint_name = tc.constraint_name
-		WHERE
-			c.table_schema = 'public'
-			AND c.table_name = $1;`
-		return c.Query(query, tableName)
+		psql := sq.StatementBuilder.PlaceholderFormat(sq.Dollar)
+
+		query, args, err := psql.Select(
+			"c.column_name",
+			"c.is_nullable",
+			"c.data_type",
+			"c.character_maximum_length",
+			"c.numeric_precision",
+			"c.numeric_scale",
+			"c.ordinal_position",
+			"tc.constraint_type AS pkey",
+		).
+			From("information_schema.columns AS c").
+			LeftJoin(
+				`information_schema.constraint_column_usage AS ccu
+					ON c.table_schema = ccu.table_schema
+						AND c.table_name = ccu.table_name
+						AND c.column_name = ccu.column_name`,
+			).
+			LeftJoin(
+				`information_schema.table_constraints AS tc
+					ON ccu.constraint_schema = tc.constraint_schema
+						AND ccu.constraint_name = tc.constraint_name`,
+			).
+			Where(
+				sq.And{
+					sq.Eq{"c.table_schema": c.schema},
+					sq.Eq{"c.table_name": tableName},
+				},
+			).
+			ToSql()
+		if err != nil {
+			return nil, nil, err
+		}
+
+		return c.Query(query, args...)
 	case "mysql":
 		query = fmt.Sprintf("DESCRIBE %s;", tableName)
 		return c.Query(query)
@@ -402,7 +432,7 @@ func (c *Client) constraints(tableName string) ([][]string, []string, error) {
 	case "postgres":
 		fallthrough
 	case "postgresql":
-		query = query.Where("tc.table_schema = 'public'")
+		query = query.Where(fmt.Sprintf("tc.table_schema = '%s'", c.schema))
 		query = query.PlaceholderFormat(sq.Dollar)
 	}
 
diff --git a/pkg/command/command.go b/pkg/command/command.go
index 37b55b6..a73979e 100644
--- a/pkg/command/command.go
+++ b/pkg/command/command.go
@@ -12,6 +12,8 @@ type Options struct {
 	Pass   string
 	DBName string
 	SSL    string
+	// PostgreSQL only.
+	Schema string
 	Limit  int
 }
 
@@ -46,5 +48,9 @@ func SetDefault(opts Options) Options {
 		opts.Port = os.Getenv("DB_PORT")
 	}
 
+	if opts.Schema == "" {
+		opts.Schema = os.Getenv("DB_SCHEMA")
+	}
+
 	return opts
 }
diff --git a/pkg/config/.dblab.yaml b/pkg/config/.dblab.yaml
index d9e592e..cf09a83 100644
--- a/pkg/config/.dblab.yaml
+++ b/pkg/config/.dblab.yaml
@@ -4,5 +4,6 @@ database:
   db: "users"
   password: "password"
   user: "postgres"
+  schema: "public"
   driver: "postgres"
 limit: 50
diff --git a/pkg/config/config.go b/pkg/config/config.go
index de1e018..043012e 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -29,6 +29,7 @@ type Config struct {
 		User     string
 		Password string
 		Driver   string `validate:"required"`
+		Schema   string
 		SSL      string `default:"disable"`
 	}
 	User   string
@@ -76,6 +77,7 @@ func Init() (command.Options, error) {
 		Pass:   cfg.Database.Password,
 		DBName: cfg.Database.DB,
 		SSL:    cfg.Database.SSL,
+		Schema: cfg.Database.Schema,
 		Limit:  cfg.Limit,
 	}
 


Test patch:
diff --git a/db/migrations/20221011024220_create_test_schema.down.sql b/db/migrations/20221011024220_create_test_schema.down.sql
new file mode 100644
index 0000000..c2d9b7d
--- /dev/null
+++ b/db/migrations/20221011024220_create_test_schema.down.sql
@@ -0,0 +1,1 @@
+DROP SCHEMA myschema CASCADE;
diff --git a/db/migrations/20221011024220_create_test_schema.up.sql b/db/migrations/20221011024220_create_test_schema.up.sql
new file mode 100644
index 0000000..138887e
--- /dev/null
+++ b/db/migrations/20221011024220_create_test_schema.up.sql
@@ -0,0 +1,1 @@
+CREATE SCHEMA myschema;
diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go
index 1716b4b..958b5f3 100644
--- a/pkg/client/client_test.go
+++ b/pkg/client/client_test.go
@@ -23,6 +23,7 @@ var (
 	host     string
 	port     string
 	name     string
+	schema   string
 )
 
 func TestMain(m *testing.M) {
@@ -32,6 +33,7 @@ func TestMain(m *testing.M) {
 	host = os.Getenv("DB_HOST")
 	port = os.Getenv("DB_PORT")
 	name = os.Getenv("DB_NAME")
+	schema = os.Getenv("DB_SCHEMA")
 
 	os.Exit(m.Run())
 }
@@ -39,7 +41,7 @@ func TestMain(m *testing.M) {
 func generateURL(driver string) string {
 	switch driver {
 	case "postgres":
-		return fmt.Sprintf("%s://%s:%s@%s:%s/%s?sslmode=mode", driver, user, password, host, port, name)
+		return fmt.Sprintf("%s://%s:%s@%s:%s/%s?sslmode=disable", driver, user, password, host, port, name)
 	case "mysql":
 		return fmt.Sprintf("%s://%s:%s@tcp(%s:%s)/%s", driver, user, password, host, port, name)
 	case "sqlite3":
@@ -84,6 +86,7 @@ func TestNewClientByUserData(t *testing.T) {
 		DBName: name,
 		SSL:    "disable",
 		Limit:  50,
+		Schema: schema,
 	}
 
 	c, err := New(opts)
@@ -106,6 +109,7 @@ func TestNewClientPing(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  50,
 	}
@@ -134,6 +138,7 @@ func TestQuery(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -159,6 +164,7 @@ func TestTableContent(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -184,6 +190,7 @@ func TestShowTables(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -208,6 +215,7 @@ func TestTableStructure(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -234,6 +242,7 @@ func TestConstraints(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -262,6 +271,7 @@ func TestIndexes(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -287,6 +297,7 @@ func TestCountTable(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -311,6 +322,7 @@ func TestMetadata(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
@@ -353,6 +365,7 @@ func TestTotalPages(t *testing.T) {
 		Host:   host,
 		Port:   port,
 		DBName: name,
+		Schema: schema,
 		SSL:    "disable",
 		Limit:  100,
 	}
diff --git a/pkg/command/command_test.go b/pkg/command/command_test.go
index 0750528..fb3f70f 100644
--- a/pkg/command/command_test.go
+++ b/pkg/command/command_test.go
@@ -36,6 +36,9 @@ func TestSetEnvValues(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	if err := os.Setenv("DB_SCHEMA", "public"); err != nil {
+		t.Fatal(err)
+	}
 	opts := Options{}
 
 	result := SetDefault(opts)
@@ -46,4 +49,5 @@ func TestSetEnvValues(t *testing.T) {
 	assert.Equal(t, "password", result.Pass)
 	assert.Equal(t, "database", result.DBName)
 	assert.Equal(t, "5432", result.Port)
+	assert.Equal(t, "public", result.Schema)
 }
diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
index 7f90dfa..ff2e1fc 100644
--- a/pkg/config/config_test.go
+++ b/pkg/config/config_test.go
@@ -17,5 +17,6 @@ func TestInit(t *testing.T) {
 	assert.Equal(t, "postgres", opts.User)
 	assert.Equal(t, "password", opts.Pass)
 	assert.Equal(t, "postgres", opts.Driver)
+	assert.Equal(t, "public", opts.Schema)
 	assert.Equal(t, 50, opts.Limit)
 }