Come impostare un tunnel SSH in Google Cloud Dataflow su un server di database esterno?

Question

Sto affrontando un problema per far funzionare la mia pipeline Apache Beam su Cloud Dataflow, con DataflowRunner.

Il primo passaggio della pipeline consiste nel connettersi a un server Postgresql esterno ospitato su una VM accessibile solo esternamente tramite SSH, porta 22 ed estrarre alcuni dati. Non posso modificare queste regole di firewall, quindi posso connettermi al server DB solo tramite tunneling SSH, noto anche come port forwarding.

Nel mio codice utilizzo la libreria python sshtunnel. Funziona perfettamente quando la pipeline viene avviata dal mio computer di sviluppo con DirectRunner :

from sshtunnel import open_tunnel

with open_tunnel(
        (user_options.ssh_tunnel_host, user_options.ssh_tunnel_port),
        ssh_username=user_options.ssh_tunnel_user,
        ssh_password=user_options.ssh_tunnel_password,
        remote_bind_address=(user_options.dbhost, user_options.dbport)
    ) as tunnel:
        with beam.Pipeline(options=pipeline_options) as p:
            (p | "Read data" >> ReadFromSQL(
                host=tunnel.local_bind_host,
                port=tunnel.local_bind_port,
                username=user_options.dbusername,
                password=user_options.dbpassword,
                database=user_options.dbname,
                wrapper=PostgresWrapper,
                query=select_query
            )
                | "Format CSV" >> DictToCSV(headers)
                | "Write CSV" >> WriteToText(user_options.export_location)
            )

Lo stesso codice, avviato con DataflowRunner all'interno di un VPC non predefinito in cui tutti gli ingressi sono negati ma nessuna restrizione in uscita e CloudNAT configurato, non riesce con questo messaggio:

psycopg2.OperationalError: impossibile connettersi al server: connessione rifiutata Il server è in esecuzione sull'host "0.0.0.0" e accetta connessioni TCP / IP sulla porta 41697? [durante l'esecuzione di "Leggi dati / Leggi"]

Quindi, ovviamente, qualcosa non va nel mio tunnel, ma non riesco a individuare esattamente cosa. Stavo iniziando a chiedermi se una configurazione diretta del tunnel SSH fosse possibile anche tramite CloudNAT, finché non ho trovato questo post sul blog:https://cloud.google.com/blog/products/gcp/guide-to-common-cloud-dataflow-use-case-patterns-part-1 affermando:

Uno dei punti di forza di Cloud Dataflow è che puoi chiamare servizi esterni per l'arricchimento dei dati. Ad esempio, puoi chiamare un micro servizio per ottenere dati aggiuntivi per un elemento. All'interno di un DoFn, chiamata al servizio (di solito eseguita tramite HTTP). Hai il pieno controllo per effettuare qualsiasi tipo di connessione che scegli, purché le regole del firewall che hai impostato all'interno del tuo progetto / rete lo consentano.

Quindi dovrebbe essere possibile creare questo tunnel! Non voglio arrendermi ma non so cosa provare dopo. Qualche idea ?

Grazie per aver letto

Answer 1

Problema risolto ! Non posso credere di aver passato due giorni interi su questo ... Stavo guardando completamente nella direzione sbagliata.

Il problema non riguardava alcune configurazioni di rete Dataflow o GCP e, per quanto ne so ...

Hai il pieno controllo per effettuare qualsiasi tipo di connessione che scegli, purché le regole del firewall che hai impostato all'interno del tuo progetto / rete lo consentano

è vero.

Il problema era ovviamente nel mio codice: solo il problema veniva rivelato solo in un ambiente distribuito. Ho commesso l'errore di aprire il tunnel dal processore della pipeline principale, invece che dai lavoratori. Quindi il tunnel SSH era attivo ma non tra i worker e il server di destinazione, solo tra la pipeline principale e il target!

Per risolvere questo problema, ho dovuto modificare il mio DoFn richiesto per racchiudere l'esecuzione della query con il tunnel:

class TunnelledSQLSourceDoFn(sql.SQLSourceDoFn):
"""Wraps SQLSourceDoFn in a ssh tunnel"""

def __init__(self, *args, **kwargs):
    self.dbport = kwargs["port"]
    self.dbhost = kwargs["host"]
    self.args = args
    self.kwargs = kwargs
    super().__init__(*args, **kwargs)

def process(self, query, *args, **kwargs):
    # Remote side of the SSH Tunnel
    remote_address = (self.dbhost, self.dbport)
    ssh_tunnel = (self.kwargs['ssh_host'], self.kwargs['ssh_port'])
    with open_tunnel(
        ssh_tunnel,
        ssh_username=self.kwargs["ssh_user"],
        ssh_password=self.kwargs["ssh_password"],
        remote_bind_address=remote_address,
        set_keepalive=10.0
    ) as tunnel:
        forwarded_port = tunnel.local_bind_port
        self.kwargs["port"] = forwarded_port
        source = sql.SQLSource(*self.args, **self.kwargs)
        sql.SQLSouceInput._build_value(source, source.runtime_params)
        logging.info("Processing - {}".format(query))
        for records, schema in source.client.read(query):
            for row in records:
                yield source.client.row_as_dict(row, schema)

come puoi vedere, ho dovuto sovrascrivere alcuni bit della libreria pysql_beam.

Infine, ogni lavoratore apre il proprio tunnel per ogni richiesta. Probabilmente è possibile ottimizzare questo comportamento ma è sufficiente per le mie esigenze.