Merge pull request #52 from citusdata/citus_dev_rebased

Introduce citus_dev

citus_dev/README.md

@@ -0,0 +1,26 @@
+# Setup
+To install dependencies run: 
+```bash
+pip install -r requirements.txt
+```
+
+Add `citus_dev` to your PATH:
+```bash
+export PATH=$PATH:<your path to citus_dev folder>
+```
+
+You can also add this to your profile:
+```bash
+echo 'export PATH=$PATH:<your path to citus_dev folder>' >>~/.profile
+```
+
+After that, you can use the citus dev tool:
+
+```bash
+citus_dev make clusterName
+```
+
+For the full command list:
+```bash
+citus_dev --help
+```

citus_dev/bash_completion

@@ -0,0 +1,52 @@
+# citus_dev completion
+
+_citus_dev()
+{
+    local cmd cur prev cmds opts clusters
+    COMPREPLY=()
+    _get_comp_words_by_ref cur prev
+    cmd="${COMP_WORDS[1]}"
+    cmds="make start stop restart"
+
+    case $prev in
+        make)
+            return 0
+            ;;
+
+        start|stop|restart)
+            clusters=$(for x in `find . -type d -name 'coordinator' | awk -F '/' '{print $2}'`; do echo ${x} ; done );
+            COMPREPLY=($( compgen -W "${clusters}" -- "${cur}" ));
+            return 0
+            ;;
+    esac
+
+    case $cmd in
+        make)
+            opts="--size --port --use-ssl --no-extension --mx"
+            if [[ ${cur} == -* ]] ; then
+                COMPREPLY=($( compgen -W "${opts}" -- "${cur}"))
+            fi
+            return 0
+            ;;
+
+        start|stop)
+            opts="--port"
+            if [[ ${cur} == -* ]] ; then
+                COMPREPLY=($( compgen -W "${opts}" -- "${cur}"))
+            fi
+            return 0
+            ;;
+
+        restart)
+            opts="--port --watch"
+            if [[ ${cur} == -* ]] ; then
+                COMPREPLY=($( compgen -W "${opts}" -- "${cur}"))
+            fi
+            return 0
+            ;;
+    esac
+
+    COMPREPLY=($( compgen -W "${cmds}" -- "${cur}"))
+    return 0
+}
+complete -F _citus_dev citus_dev

citus_dev/citus_dev

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""citus_dev
+
+Usage:
+  citus_dev make <name> [--size=<count>] [--port=<port>] [--use-ssl] [--no-extension] [--mx] [--destroy] [--init-with=<sql_file>]
+  citus_dev restart <name> [--watch] [--port=<port>]
+  citus_dev (start|stop) <name> [--port=<port>]
+
+Options:
+  --size=<count>           Number of workers to create when 0 the coordinator will be added as a worker [default: 2]
+  --port=<port>            Port number to use for the coordinator. All workers take subsequent numbers [default: 9700]
+  --watch                  Watch for changes to the citus plugin and restart the cluster when the plugin updates
+  --use-ssl                Create the cluster with ssl enabled
+  --no-extension           Do not create the extension while creating the nodes
+  --mx                     Enable metadata sync for all workers
+  --destroy                Destroy any old cluster with the same name
+  --init-with=<sql_file>   A SQL script to run after creation of the cluster to set up any necessary tables and data
+
+"""
+from docopt import docopt
+from subprocess import call
+from subprocess import Popen, PIPE
+import os
+import sys
+
+
+def createNodeCommands(clustername, role, index=None, usessl=False, mx=False):
+    cs = []
+
+    nodename = role
+    if index != None:
+        nodename += "%d" % index
+
+    dir = "%s/%s" % (clustername, nodename)
+    cs.append("initdb -D %s" % dir)
+    cs.append("echo \"shared_preload_libraries = 'citus'\" >> %s/postgresql.conf" % dir)
+    cs.append('echo "wal_level = logical" >> %s/postgresql.conf' % dir)
+
+    if usessl:
+        cs.append('echo "ssl = on" >> %s/postgresql.conf' % dir)
+        cs.append(
+            "echo \"citus.node_conninfo = 'sslmode=require'\" >> %s/postgresql.conf"
+            % dir
+        )
+        cs.append(
+            "openssl req -new -x509 -days 365 -nodes -text -out %s/server.crt -keyout %s/server.key -subj '/CN=%s'"
+            % (dir, dir, nodename)
+        )
+        cs.append("chmod 0600 %s/server.key" % dir)
+
+    if mx:
+        cs.append(
+            "echo \"citus.replication_model = 'streaming'\" >> %s/postgresql.conf" % dir
+        )
+
+    return cs
+
+
+def main(arguments):
+    print(arguments)
+    if arguments["make"]:
+        cs = []
+        if arguments['--destroy']:
+            name = arguments["<name>"]
+            for role in getRoles(name):
+                cs.append("pg_ctl stop -D %s/%s" % (name, role))
+            cs.append('rm -rf %s' % (name))
+
+        cs += createNodeCommands(
+            arguments["<name>"],
+            "coordinator",
+            usessl=arguments["--use-ssl"],
+            mx=arguments["--mx"],
+        )
+
+        size = int(arguments["--size"])
+
+        for i in range(size):
+            cs += createNodeCommands(
+                arguments["<name>"],
+                "worker",
+                i,
+                usessl=arguments["--use-ssl"],
+                mx=arguments["--mx"],
+            )
+
+        port = int(arguments["--port"])
+
+        cport = port
+        role = "coordinator"
+        cs.append(
+            'pg_ctl -D %s/%s -o "-p %d" -l %s_logfile start'
+            % (arguments["<name>"], role, cport, role)
+        )
+        port += 1
+
+        for i in range(size):
+            role = "worker%d" % i
+            cs.append(
+                'pg_ctl start -D %s/%s -o "-p %d" -l %s_logfile'
+                % (arguments["<name>"], role, port, role)
+            )
+            port += 1
+        port = cport
+
+
+        if not arguments["--no-extension"]:
+            for i in range(size + 1):
+                cs.append('psql -p %d -c "CREATE EXTENSION citus;"' % (port + i))
+
+            # If the cluster size is 0 we add the coordinator as the only node, otherwise we will add all other nodes
+            if size == 0:
+                cs.append(
+                    "psql -p %d -c \"SELECT * from master_add_node('localhost', %d);\""
+                    % (port, port)
+                )
+            else:
+                for i in range(size):
+                    cs.append(
+                        "psql -p %d -c \"SELECT * from master_add_node('localhost', %d);\""
+                        % (port, port + 1 + i)
+                    )
+                    if arguments["--mx"]:
+                        cs.append(
+                            "psql -p %d -c \"SELECT start_metadata_sync_to_node('localhost', %d);\""
+                            % (port, port + 1 + i)
+                        )
+
+            cs.append(
+                'psql -p %d -c "SELECT * from master_get_active_worker_nodes();"'
+                % (port)
+            )
+        if arguments['--init-with']:
+            cs.append('psql -p %d -f %s' % (cport, arguments['--init-with']))
+
+        for c in cs:
+            print(c)
+            os.system(c)
+            print("")
+
+    elif arguments["stop"]:
+        cs = []
+        name = arguments["<name>"]
+        for role in getRoles(name):
+            cs.append("pg_ctl stop -D %s/%s" % (name, role))
+
+        for c in cs:
+            print(c)
+            os.system(c)
+            print("")
+
+    elif arguments["start"]:
+        cs = []
+        name = arguments["<name>"]
+        port = int(arguments["--port"])
+        cport = port
+        for role in getRoles(name):
+            cs.append(
+                'pg_ctl start -D %s/%s -o "-p %d" -l %s_logfile'
+                % (name, role, cport, role)
+            )
+            cport += 1
+
+        for c in cs:
+            print(c)
+            os.system(c)
+            print("")
+
+    elif arguments["restart"]:
+        cs = []
+        name = arguments["<name>"]
+        port = int(arguments["--port"])
+        if arguments["--watch"]:
+            cs.append(
+                "fswatch -0 '%s' | xargs -0 -n 1 -I{} citus_dev restart %s --port=%d"
+                % (citus_so(), name, port)
+            )
+
+        else:
+            cport = port
+            for role in getRoles(name):
+                cs.append(
+                    'pg_ctl restart -D %s/%s -o "-p %d" -l %s_logfile'
+                    % (name, role, cport, role)
+                )
+                cport += 1
+
+        for c in cs:
+            print(c)
+            os.system(c)
+            print("")
+
+    else:
+        print("unknown command")
+        exit(1)
+
+
+def getRoles(name):
+    roles = [f.name for f in os.scandir(name) if f.is_dir()]
+    roles.sort()
+    return roles
+
+
+def pg_libdir():
+    process = Popen(["pg_config"], stdout=PIPE)
+    (output, err) = process.communicate()
+    exit_code = process.wait()
+
+    output = str(output)
+
+    for line in output.split("\\n"):
+        if line.startswith("LIBDIR"):
+            return line.split("=", 1)[1].strip()
+
+    raise Exception("can't find postgres lib dir")
+
+
+def citus_so():
+    return pg_libdir() + "/citus.so"
+
+
+if __name__ == "__main__":
+    print(sys.argv)
+    main(docopt(__doc__, version="citus_dev"))

citus_dev/requirements.txt

@@ -0,0 +1 @@
+docopt==0.6.2