scripts/analytics/db_connect.R

# Shared connection helper (SSH tunnel + RPostgres), mirrors analytics scripts.
zabota_connect <- function() {
  db_host <- Sys.getenv("DB_HOST", "localhost")
  db_port <- as.integer(Sys.getenv("DB_PORT", "5432"))
  db_name <- Sys.getenv("DB_NAME", "superset_data")
  db_user <- Sys.getenv("DB_USER", "superset_data")
  db_password <- Sys.getenv("DB_PASSWORD", "")
  db_sslmode <- Sys.getenv("DB_SSLMODE", "prefer")
  db_gssencmode <- Sys.getenv("DB_GSSENCMODE", "disable")
  use_ssh_tunnel <- tolower(Sys.getenv("USE_SSH_TUNNEL", "false")) %in% c("1", "true", "yes")

  wait_for_local_port <- function(host, port, timeout_sec = 8) {
    deadline <- Sys.time() + timeout_sec
    repeat {
      cn <- try(suppressWarnings(socketConnection(host, port, open = "r+", blocking = TRUE, timeout = 1)), silent = TRUE)
      if (!inherits(cn, "try-error")) { close(cn); return(TRUE) }
      if (Sys.time() > deadline) return(FALSE)
      Sys.sleep(0.3)
    }
  }

  if (use_ssh_tunnel) {
    ssh_host <- Sys.getenv("SSH_HOST", db_host)
    ssh_user <- Sys.getenv("SSH_USER", "")
    ssh_key <- Sys.getenv("SSH_KEY", "~/.ssh/id_rsa")
    ssh_local_port <- as.integer(Sys.getenv("SSH_LOCAL_PORT", "5433"))
    ssh_remote_db_host <- Sys.getenv("SSH_REMOTE_DB_HOST", "localhost")
    ssh_remote_db_port <- as.integer(Sys.getenv("SSH_REMOTE_DB_PORT", as.character(db_port)))
    args <- c("-f", "-N", "-o", "ExitOnForwardFailure=yes",
              "-L", paste0(ssh_local_port, ":", ssh_remote_db_host, ":", ssh_remote_db_port),
              "-i", path.expand(ssh_key), paste0(ssh_user, "@", ssh_host))
    out <- tryCatch(system2("ssh", args, stdout = TRUE, stderr = TRUE),
                    warning = function(w) character(), error = function(e) conditionMessage(e))
    st <- attr(out, "status"); st <- if (is.null(st)) 0L else st
    txt <- paste(out, collapse = "\n")
    if (st != 0L && !grepl("Address already in use", txt, fixed = TRUE)) stop("SSH tunnel failed:\n", txt)
    db_host <- "localhost"; db_port <- ssh_local_port
    if (!wait_for_local_port("127.0.0.1", db_port)) stop("tunnel port unreachable")
  }

  DBI::dbConnect(RPostgres::Postgres(), host = db_host, port = db_port, dbname = db_name,
                 user = db_user, password = db_password, sslmode = db_sslmode, gssencmode = db_gssencmode)
}