ssh tunnel for postgres normalization

.github/workflows/publish-command.yml

@@ -86,6 +86,8 @@ jobs:
           BRAINTREE_TEST_CREDS: ${{ secrets.BRAINTREE_TEST_CREDS }}
           CART_TEST_CREDS: ${{ secrets.CART_TEST_CREDS }}
           CHARGEBEE_INTEGRATION_TEST_CREDS: ${{ secrets.CHARGEBEE_INTEGRATION_TEST_CREDS }}
+          DESTINATION_POSTGRES_SSH_KEY_TEST_CREDS: ${{ secrets.DESTINATION_POSTGRES_SSH_KEY_TEST_CREDS }}
+          DESTINATION_POSTGRES_SSH_PWD_TEST_CREDS: ${{ secrets.DESTINATION_POSTGRES_SSH_PWD_TEST_CREDS }}
           DESTINATION_PUBSUB_TEST_CREDS: ${{ secrets.DESTINATION_PUBSUB_TEST_CREDS }}
           DESTINATION_KEEN_TEST_CREDS: ${{ secrets.DESTINATION_KEEN_TEST_CREDS }}
           DESTINATION_KVDB_TEST_CREDS: ${{ secrets.DESTINATION_KVDB_TEST_CREDS }}

.github/workflows/test-command.yml

@@ -86,6 +86,8 @@ jobs:
           BRAINTREE_TEST_CREDS: ${{ secrets.BRAINTREE_TEST_CREDS }}
           CART_TEST_CREDS: ${{ secrets.CART_TEST_CREDS }}
           CHARGEBEE_INTEGRATION_TEST_CREDS: ${{ secrets.CHARGEBEE_INTEGRATION_TEST_CREDS }}
+          DESTINATION_POSTGRES_SSH_KEY_TEST_CREDS: ${{ secrets.DESTINATION_POSTGRES_SSH_KEY_TEST_CREDS }}
+          DESTINATION_POSTGRES_SSH_PWD_TEST_CREDS: ${{ secrets.DESTINATION_POSTGRES_SSH_PWD_TEST_CREDS }}
           DESTINATION_PUBSUB_TEST_CREDS: ${{ secrets.DESTINATION_PUBSUB_TEST_CREDS }}
           DESTINATION_KEEN_TEST_CREDS: ${{ secrets.DESTINATION_KEEN_TEST_CREDS }}
           DESTINATION_KVDB_TEST_CREDS: ${{ secrets.DESTINATION_KVDB_TEST_CREDS }}

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/AirbyteMessageConsumer.java

@@ -24,6 +24,7 @@
 
 package io.airbyte.integrations.base;
 
+import io.airbyte.commons.concurrency.VoidCallable;
 import io.airbyte.commons.functional.CheckedConsumer;
 import io.airbyte.protocol.models.AirbyteMessage;
 
@@ -56,4 +57,29 @@ public interface AirbyteMessageConsumer extends CheckedConsumer<AirbyteMessage,
   @Override
   void close() throws Exception;
 
+  /**
+   * Append a function to be called on {@link AirbyteMessageConsumer#close}.
+   */
+  static AirbyteMessageConsumer appendOnClose(final AirbyteMessageConsumer consumer, final VoidCallable voidCallable) {
+    return new AirbyteMessageConsumer() {
+
+      @Override
+      public void start() throws Exception {
+        consumer.start();
+      }
+
+      @Override
+      public void accept(final AirbyteMessage message) throws Exception {
+        consumer.accept(message);
+      }
+
+      @Override
+      public void close() throws Exception {
+        consumer.close();
+        voidCallable.call();
+      }
+
+    };
+  }
+
 }

airbyte-integrations/bases/base-java/src/main/java/io/airbyte/integrations/base/ssh/SshWrappedDestination.java

@@ -0,0 +1,81 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2020 Airbyte
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package io.airbyte.integrations.base.ssh;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import io.airbyte.commons.json.Jsons;
+import io.airbyte.commons.resources.MoreResources;
+import io.airbyte.integrations.base.AirbyteMessageConsumer;
+import io.airbyte.integrations.base.Destination;
+import io.airbyte.protocol.models.AirbyteConnectionStatus;
+import io.airbyte.protocol.models.AirbyteMessage;
+import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+import io.airbyte.protocol.models.ConnectorSpecification;
+import java.util.List;
+import java.util.function.Consumer;
+
+/**
+ * Decorates a Destination with an SSH Tunnel using the standard configuration that Airbyte uses for
+ * configuring SSH.
+ */
+public class SshWrappedDestination implements Destination {
+
+  private final Destination delegate;
+  private final List<String> hostKey;
+  private final List<String> portKey;
+
+  public SshWrappedDestination(final Destination delegate,
+                               final List<String> hostKey,
+                               final List<String> portKey) {
+    this.delegate = delegate;
+    this.hostKey = hostKey;
+    this.portKey = portKey;
+  }
+
+  @Override
+  public ConnectorSpecification spec() throws Exception {
+    // inject the standard ssh configuration into the spec.
+    final ConnectorSpecification originalSpec = delegate.spec();
+    final ObjectNode propNode = (ObjectNode) originalSpec.getConnectionSpecification().get("properties");
+    propNode.set("tunnel_method", Jsons.deserialize(MoreResources.readResource("ssh-tunnel-spec.json")));
+    return originalSpec;
+  }
+
+  @Override
+  public AirbyteConnectionStatus check(final JsonNode config) throws Exception {
+    return SshTunnel.sshWrap(config, hostKey, portKey, delegate::check);
+  }
+
+  @Override
+  public AirbyteMessageConsumer getConsumer(final JsonNode config,
+                                            final ConfiguredAirbyteCatalog catalog,
+                                            final Consumer<AirbyteMessage> outputRecordCollector)
+      throws Exception {
+    final SshTunnel tunnel = SshTunnel.getInstance(config, hostKey, portKey);
+    return AirbyteMessageConsumer.appendOnClose(delegate.getConsumer(tunnel.getConfigInTunnel(), catalog, outputRecordCollector), tunnel::close);
+  }
+
+}

airbyte-integrations/bases/base-normalization/.dockerignore

@@ -1,6 +1,7 @@
 *
 !Dockerfile
 !entrypoint.sh
+!ssh/sshtunneling.sh
 !setup.py
 !normalization
 !dbt-project-template

airbyte-integrations/bases/base-normalization/Dockerfile

@@ -1,8 +1,11 @@
 FROM fishtownanalytics/dbt:0.19.0
 COPY --from=airbyte/base-airbyte-protocol-python:0.1.1 /airbyte /airbyte
 
+RUN apt-get update && apt-get install -y jq sshpass
+
 WORKDIR /airbyte
 COPY entrypoint.sh .
+COPY ssh/sshtunneling.sh .
 
 WORKDIR /airbyte/normalization_code
 COPY normalization ./normalization
@@ -24,5 +27,5 @@ WORKDIR /airbyte
 ENV AIRBYTE_ENTRYPOINT "/airbyte/entrypoint.sh"
 ENTRYPOINT ["/airbyte/entrypoint.sh"]
 
-LABEL io.airbyte.version=0.1.42
+LABEL io.airbyte.version=0.1.43
 LABEL io.airbyte.name=airbyte/normalization

airbyte-integrations/bases/base-normalization/entrypoint.sh

@@ -94,8 +94,11 @@ function main() {
   case "$CMD" in
   run)
     configuredbt
+    . /airbyte/sshtunneling.sh
+    openssh $CONFIG_FILE "${PROJECT_DIR}/localsshport.json"
     # Run dbt to compile and execute the generated normalization models
     dbt run --profiles-dir "${PROJECT_DIR}" --project-dir "${PROJECT_DIR}"
+    closessh
     ;;
   configure-dbt)
     configuredbt

airbyte-integrations/bases/base-normalization/normalization/transform_config/transform.py

@@ -27,6 +27,7 @@
 import json
 import os
 import pkgutil
+import socket
 from enum import Enum
 from typing import Any, Dict
 
@@ -48,6 +49,8 @@ def run(self, args):
         integration_type = inputs["integration_type"]
         transformed_config = self.transform(integration_type, original_config)
         self.write_yaml_config(inputs["output_path"], transformed_config)
+        if self.is_ssh_tunnelling(original_config):
+            self.write_ssh_port(inputs["output_path"], self.pick_a_port())
 
     @staticmethod
     def parse(args):
@@ -86,6 +89,43 @@ def transform(self, integration_type: DestinationType, config: Dict[str, Any]):
 
         return base_profile
 
+    @staticmethod
+    def is_ssh_tunnelling(config: Dict[str, Any]) -> bool:
+        tunnel_methods = ["SSH_KEY_AUTH", "SSH_PASSWORD_AUTH"]
+        if (
+            "tunnel_method" in config.keys() and 
+            "tunnel_method" in config["tunnel_method"] and 
+            config["tunnel_method"]["tunnel_method"].upper() in tunnel_methods
+        ):
+            return True
+        else:
+            return False
+
+    @staticmethod
+    def is_port_free(port: int) -> bool:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.bind(("localhost", port))
+            except Exception as e:
+                print(f"port {port} unsuitable: {e}")
+                return False
+            else:
+                print(f"port {port} is free")
+                return True
+
+    @staticmethod
+    def pick_a_port() -> int:
+        """
+        This function finds a free port, starting with 50001 and adding 1 until we find an open port.
+        """
+        port_to_check = 50001  # just past start of dynamic port range (49152:65535)
+        while not TransformConfig.is_port_free(port_to_check):
+            port_to_check += 1
+            # error if we somehow hit end of port range
+            if port_to_check > 65535:
+                raise RuntimeError("Couldn't find a free port to use.")
+        return port_to_check
+
     @staticmethod
     def transform_bigquery(config: Dict[str, Any]):
         print("transform_bigquery")
@@ -108,13 +148,22 @@ def transform_bigquery(config: Dict[str, Any]):
     @staticmethod
     def transform_postgres(config: Dict[str, Any]):
         print("transform_postgres")
+
+        # set port & host correctly depending on whether we're ssh tunnelling
+        # TODO: this should be a separate function to stay DRY when adding support for other destinations
+        port = config["port"]
+        host = config["host"]
+        if TransformConfig.is_ssh_tunnelling(config):
+            port = TransformConfig.pick_a_port()
+            host = "localhost"
+
         # https://docs.getdbt.com/reference/warehouse-profiles/postgres-profile
         dbt_config = {
             "type": "postgres",
-            "host": config["host"],
+            "host": host,
             "user": config["username"],
             "pass": config.get("password", ""),
-            "port": config["port"],
+            "port": port,
             "dbname": config["database"],
             "schema": config["schema"],
             "threads": 32,
@@ -191,6 +240,19 @@ def write_yaml_config(output_path: str, config: Dict[str, Any]):
         with open(os.path.join(output_path, "profiles.yml"), "w") as fh:
             fh.write(yaml.dump(config))
 
+    @staticmethod
+    def write_ssh_port(output_path: str, port: int):
+        """
+        This function writes a small json file with content like {"port":xyz}
+        This is being used only when ssh tunneling.
+        We do this because we need to decide on and save this port number into our dbt config
+        and then use that same port in sshtunneling.sh when opening the tunnel.
+        """
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+        with open(os.path.join(output_path, "localsshport.json"), "w") as fh:
+            json.dump({"port": port}, fh)
+
 
 def main(args=None):
     TransformConfig().run(args)

airbyte-integrations/bases/base-normalization/ssh/sshtunneling.sh

@@ -0,0 +1,62 @@
+# This function opens an ssh tunnel if required using values provided in config.
+# Requires two arguments,
+  # path to config file ($1)
+  # path to file containing local port to use ($2)
+function openssh() {
+  # check if jq is missing, and if so try to install it..
+  # this is janky but for custom dbt transform we can't be sure jq is installed as using user docker image
+  if ! command -v jq &> /dev/null ; then
+    echo "CRITICAL: jq not installed... attempting to install on the fly but will fail if unable."
+    { apt-get update && apt-get install -y jq; } ||
+    apk --update add jq ||
+    { yum install epel-release -y && yum install jq -y; } ||
+    { dnf install epel-release -y && dnf install jq -y; } || exit 1
+  fi
+  tunnel_method=$(cat $1 | jq -r '.tunnel_method.tunnel_method' | tr '[:lower:]' '[:upper:]')
+  tunnel_username=$(cat $1 | jq -r '.tunnel_method.tunnel_user')
+  tunnel_db_host=$(cat $1 | jq -r '.host')
+  tunnel_db_port=$(cat $1 | jq -r '.port')
+  tunnel_host=$(cat $1 | jq -r '.tunnel_method.tunnel_host')
+  tunnel_local_port=$(cat $2 | jq -r '.port')
+  # set a path for a control socket, allowing us to close this specific ssh connection when desired
+  tmpcontrolsocket="/tmp/sshsocket${tunnel_db_remote_port}-${RANDOM}"
+  if [[ ${tunnel_method} = "SSH_KEY_AUTH" ]] ; then
+    echo "Detected tunnel method SSH_KEY_AUTH for normalization"
+    # create a temporary file to hold ssh key and trap to delete on EXIT
+    trap 'rm -f "$tmpkeyfile"' EXIT
+    tmpkeyfile=$(mktemp /tmp/xyzfile.XXXXXXXXXXX) || exit 1
+    echo "$(cat $1 | jq -r '.tunnel_method.ssh_key')" > $tmpkeyfile
+    # -f=background  -N=no remote command  -M=master mode  StrictHostKeyChecking=no auto-adds host
+    echo "Running: ssh -f -N -M -o StrictHostKeyChecking=no -S {control socket} -i {key file} -l ${tunnel_username} -L ${tunnel_local_port}:${tunnel_db_host}:${tunnel_db_port} ${tunnel_host}"
+    ssh -f -N -M -o StrictHostKeyChecking=no -S $tmpcontrolsocket -i $tmpkeyfile -l ${tunnel_username} -L ${tunnel_local_port}:${tunnel_db_host}:${tunnel_db_port} ${tunnel_host} &&
+    sshopen="yes" &&
+    echo "ssh tunnel opened"
+    rm -f $tmpkeyfile
+  elif [[ ${tunnel_method} = "SSH_PASSWORD_AUTH" ]] ; then
+    echo "Detected tunnel method SSH_PASSWORD_AUTH for normalization"
+    if ! command -v sshpass &> /dev/null ; then
+      echo "CRITICAL: sshpass not installed... attempting to install on the fly but will fail if unable."
+      { apt-get update && apt-get install -y sshpass; } ||
+      { apk add --update openssh && apk --update add sshpass; } ||
+      { yum install epel-release -y && yum install sshpass -y; } ||
+      { dnf install epel-release -y && dnf install sshpass -y; } || exit 1
+    fi
+    # put ssh password in env var for use in sshpass. Better than directly passing with -p
+    export SSHPASS=$(cat $1 | jq -r '.tunnel_method.tunnel_user_password')
+    echo "Running: sshpass -e ssh -f -N -M -o StrictHostKeyChecking=no -S {control socket} -l ${tunnel_username} -L ${tunnel_local_port}:${tunnel_db_host}:${tunnel_db_port} ${tunnel_host}"
+    sshpass -e ssh -f -N -M -o StrictHostKeyChecking=no -S $tmpcontrolsocket -l ${tunnel_username} -L ${tunnel_local_port}:${tunnel_db_host}:${tunnel_db_port} ${tunnel_host} &&
+    sshopen="yes" &&
+    echo "ssh tunnel opened"
+  fi
+}
+
+# This function checks if $sshopen variable has been set and if so, closes the ssh open via $tmpcontrolsocket
+# This only works after calling openssh()
+function closessh() {
+  # $sshopen $tmpcontrolsocket comes from openssh() function
+  if [ ! -z "$sshopen" ] ; then
+    ssh -S $tmpcontrolsocket -O exit ${tunnel_host} &&
+    echo "closed ssh tunnel"
+    trap 'rm -f "$tmpcontrolsocket"' EXIT
+  fi
+}