first commit

.gitignore

@@ -0,0 +1 @@
+.pgsql/data

.pgsql/run.sh

@@ -0,0 +1,25 @@
+#!/bin/sh
+# shellcheck shell=sh disable=SC2129
+set -e
+
+# Reset the database
+pg_ctl -D .pgsql/data -w stop -m fast || true
+rm -rf .pgsql/data
+mkdir -p .pgsql/data
+pg_ctl -D .pgsql/data init;
+echo "log_statement=all" >> .pgsql/data/postgresql.conf
+echo "max_connections=100" >> .pgsql/data/postgresql.conf
+echo "log_connections=false" >> .pgsql/data/postgresql.conf
+echo "log_disconnections=false" >> .pgsql/data/postgresql.conf
+echo "log_duration=true" >> .pgsql/data/postgresql.conf
+echo "timezone=utc" >> .pgsql/data/postgresql.conf
+echo "unix_socket_directories='/tmp'" >>.pgsql/data/postgresql.conf
+
+# Configure the databases
+pg_ctl -D .pgsql/data -w start
+LANG=en_US.UTF-8 createdb -h localhost -p 5432
+.pgsql/testdb.sh
+pg_ctl -D .pgsql/data -w stop -m fast
+
+# Start the server
+postgres -D ./.pgsql/data

.pgsql/testdb.sh

@@ -0,0 +1,18 @@
+#!/bin/sh
+# shellcheck shell=sh
+set -e
+
+
+psql -v ON_ERROR_STOP=1 --no-password --no-psqlrc "postgresql://$POSTGRES_USER@localhost:5432/$POSTGRES_DB" <<-EOSQL
+  CREATE USER test_source WITH PASSWORD '';
+  CREATE DATABASE test_source;
+  GRANT ALL PRIVILEGES ON DATABASE test_source TO test_source;
+  ALTER DATABASE test_source OWNER TO test_source;
+EOSQL
+
+psql -v ON_ERROR_STOP=1 --no-password --no-psqlrc "postgresql://$POSTGRES_USER@localhost:5432/$POSTGRES_DB" <<-EOSQL
+  CREATE USER test_target WITH PASSWORD '';
+  CREATE DATABASE test_target;
+  GRANT ALL PRIVILEGES ON DATABASE test_target TO test_target;
+  ALTER DATABASE test_target OWNER TO test_target;
+EOSQL

Makefile

@@ -0,0 +1,17 @@
+PACKAGE=niteo.co/subsetter
+
+.PHONY: run
+run:
+	go run "${PACKAGE}/$(filter-out $@,$(MAKECMDGOALS))" 
+
+.PHONY: up
+up:
+	process-compose up -t=false -p=0
+
+.PHONY: is-postgres-running
+is-postgres-running:
+	@(pg_isready -h localhost) || (echo "# ==> Startis postgres by running 'make up'" && exit 2)
+
+.PHONY: pgweb
+pgweb:is-postgres-running
+	@pgweb --url "postgres://test_source@localhost:5432/test_source?sslmode=disable"

README.md

@@ -0,0 +1,2 @@
+# pg-subsetter
+# pg-subsetter

cli/main.go

@@ -0,0 +1,17 @@
+package main
+
+import (
+	"flag"
+
+	"github.com/rs/zerolog/log"
+)
+
+var src = flag.String("src", "", "Source DSN")
+var dst = flag.String("dst", "", "Destination DSN")
+var fraction = flag.Float64("f", 1.0, "Fraction of rows to copy")
+
+func main() {
+	flag.Parse()
+	log.Info().Msg("Starting")
+
+}

flake.nix

@@ -0,0 +1,45 @@
+{
+  nixConfig = {
+    allowed-users = [ "@wheel" "@staff" ]; # allow compiling on every device/machine
+  };
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/release-23.05";
+    flake-parts.url = "github:hercules-ci/flake-parts";
+  };
+  outputs = inputs@{ self, nixpkgs, flake-parts, ... }:
+    flake-parts.lib.mkFlake { inherit inputs; } {
+
+      systems = nixpkgs.lib.systems.flakeExposed;
+      imports = [
+        inputs.flake-parts.flakeModules.easyOverlay
+      ];
+
+      perSystem = { config, self', inputs', pkgs, system, ... }:
+        let
+
+          # dev env without compile tools
+          stdenvMinimal = pkgs.stdenvNoCC.override {
+            cc = null;
+            preHook = "";
+            allowedRequisites = null;
+            initialPath = pkgs.lib.filter
+              (a: pkgs.lib.hasPrefix "coreutils" a.name)
+              pkgs.stdenvNoCC.initialPath;
+            extraNativeBuildInputs = [ ];
+          };
+        in
+        {
+          devShells.default = pkgs.mkShell {
+            stdenv = stdenvMinimal;
+            packages = with pkgs; [
+              go
+              postgresql
+              process-compose
+              shellcheck
+              nixpkgs-fmt
+              pgweb
+            ];
+          };
+        };
+    };
+}

go.mod

@@ -0,0 +1,21 @@
+module niteo.co/subsetter
+
+go 1.20
+
+require github.com/rs/zerolog v1.30.0
+
+require (
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+)
+
+require (
+	github.com/jackc/pgx v3.6.2+incompatible // indirect
+	github.com/jackc/pgx/v5 v5.4.3
+	github.com/mattn/go-colorable v0.1.12 // indirect
+	github.com/mattn/go-isatty v0.0.14 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	golang.org/x/crypto v0.12.0 // indirect
+	golang.org/x/sys v0.11.0 // indirect
+	golang.org/x/text v0.12.0 // indirect
+)

go.sum

@@ -0,0 +1,34 @@
+github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx v3.6.2+incompatible h1:2zP5OD7kiyR3xzRYMhOcXVvkDZsImVXfj+yIyTQf3/o=
+github.com/jackc/pgx v3.6.2+incompatible/go.mod h1:0ZGrqGqkRlliWnWB4zKnWtjbSWbGkVEFm4TeybAXq+I=
+github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
+github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
+github.com/mattn/go-colorable v0.1.12 h1:jF+Du6AlPIjs2BiUiQlKOX0rt3SujHxPnksPKZbaA40=
+github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
+github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
+github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.30.0 h1:SymVODrcRsaRaSInD9yQtKbtWqwsfoPcRff/oRXLj4c=
+github.com/rs/zerolog v1.30.0/go.mod h1:/tk+P47gFdPXq4QYjvCmT5/Gsug2nagsFWBWhAiSi1w=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/crypto v0.12.0 h1:tFM/ta59kqch6LlvYnPa0yx5a83cL2nHflFhYKvv9Yk=
+golang.org/x/crypto v0.12.0/go.mod h1:NF0Gs7EO5K4qLn+Ylc+fih8BSTeIjAP05siRnAh98yw=
+golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.11.0 h1:eG7RXZHdqOJ1i+0lgLgCpSXAp6M3LYlAo6osgSi0xOM=
+golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/text v0.12.0 h1:k+n5B8goJNdU7hSvEtMUz3d1Q6D/XW4COJSJR6fN0mc=
+golang.org/x/text v0.12.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

process-compose.yml

@@ -0,0 +1,10 @@
+version: "0.5"
+
+processes:
+  postgresql:
+    command: ".pgsql/run.sh"
+    availability:
+      restart: "always"
+    environment:
+      - "PGDATA=./.pgsql/data"
+      - "PGHOST=./.pgsql"

subsetter/db_test.go

@@ -0,0 +1,45 @@
+package subsetter
+
+import (
+	"context"
+	"os"
+
+	"github.com/jackc/pgx/v5"
+)
+
+func getTestConnection() *pgx.Conn {
+	DATABASE_URL := os.Getenv("DATABASE_URL")
+	if DATABASE_URL == "" {
+		DATABASE_URL = "postgres://test_source@localhost:5432/test_source?sslmode=disable"
+	}
+
+	conn, err := pgx.Connect(context.Background(), DATABASE_URL)
+	if err != nil {
+		panic(err)
+	}
+	return conn
+}
+
+func populateTests(conn *pgx.Conn) {
+	conn.Exec(context.Background(), `
+		CREATE TABLE simple (
+			id UUID PRIMARY KEY,
+			text TEXT
+		);
+		
+		CREATE TABLE relation (
+			id UUID PRIMARY KEY,
+			simple_id UUID
+		);
+
+		ALTER TABLE relation ADD CONSTRAINT relation_simple_fk FOREIGN KEY (simple_id) REFERENCES simple(id);
+	`)
+}
+
+func clearPopulateTests(conn *pgx.Conn) {
+	conn.Exec(context.Background(), `
+		ALTER TABLE relation DROP CONSTRAINT relation_simple_fk;
+		DROP TABLE simple;
+		DROP TABLE relation;
+	`)
+}

subsetter/query.go

@@ -0,0 +1,42 @@
+package subsetter
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type Table struct {
+	Name string
+	Rows int
+}
+
+func GetTables(conn *pgx.Conn) (tables []string, err error) {
+	q := `SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';`
+	rows, err := conn.Query(context.Background(), q)
+	for rows.Next() {
+		var name string
+		rows.Scan(&name)
+		tables = append(tables, name)
+	}
+	rows.Close()
+	return
+}
+
+func GetTablesWithRows(conn *pgx.Conn) (tables []Table, err error) {
+	q := `SELECT relname, reltuples::int FROM pg_class,information_schema.tables WHERE table_schema = 'public' AND relname = table_name;`
+	rows, err := conn.Query(context.Background(), q)
+	for rows.Next() {
+		var table Table
+		rows.Scan(&table.Name, &table.Rows)
+
+		// fix for tables with no rows
+		if table.Rows == -1 {
+			table.Rows = 0
+		}
+		tables = append(tables, table)
+	}
+	rows.Close()
+
+	return
+}