From f61b24cc0b04231319e37bb018d596044a33a777 Mon Sep 17 00:00:00 2001
From: boB Rudis <bob@rud.is>
Date: Sat, 16 Feb 2019 07:35:55 -0500
Subject: [PATCH] dplyr bits working

---
 NAMESPACE                                |  3 ++
 R/z-dbGetQuery.R                         | 84 ++++++++++++++++++++++++--------
 man/fetch-AthenaResult-numeric-method.Rd | 24 +++++++++
 3 files changed, 91 insertions(+), 20 deletions(-)
 create mode 100644 man/fetch-AthenaResult-numeric-method.Rd

diff --git a/NAMESPACE b/NAMESPACE
index e85874e..69e7ba8 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -11,14 +11,17 @@ exportClasses(AthenaDriver)
 exportClasses(AthenaResult)
 exportMethods(dbConnect)
 exportMethods(dbExistsTable)
+exportMethods(dbGetInfo)
 exportMethods(dbGetQuery)
 exportMethods(dbListFields)
 exportMethods(dbListTables)
 exportMethods(dbReadTable)
 exportMethods(dbSendQuery)
+exportMethods(fetch)
 import(DBI)
 import(RJDBC)
 import(bit64)
+import(dbplyr)
 importFrom(aws.signature,read_credentials)
 importFrom(aws.signature,use_credentials)
 importFrom(rJava,.jcall)
diff --git a/R/z-dbGetQuery.R b/R/z-dbGetQuery.R
index 5440275..23fa372 100644
--- a/R/z-dbGetQuery.R
+++ b/R/z-dbGetQuery.R
@@ -23,45 +23,89 @@ list(
   "1111" = as.character # OTHER
 ) -> .jdbc_converters
 
-#' AthenaJDBC
-#'
-#' @param conn Athena connection
-#' @param statement SQL statement
-#' @param ... unused
-#' @importFrom rJava .jcall
 #' @export
-setMethod(
+#' @keywords internal
+setMethod("dbGetInfo", "AthenaDriver", def=function(dbObj, ...)
+  list(
+    name = "AthenaJDBC",
+    driver_version = list.files(system.file("java", package="metis.lite"), "jar$")[1],
+    package_version = utils::packageVersion("metis.lite")
+  )
+)
 
-  "dbGetQuery",
-  signature(conn="AthenaConnection", statement="character"),
-
-  definition = function(conn, statement, type_convert=FALSE, ...) {
+#' @export
+#' @keywords internal
+setMethod("dbGetInfo", "AthenaConnection", def=function(dbObj, ...)
+  list(
+    name = "AthenaJDBC",
+    driver_version = list.files(system.file("java", package="metis.lite"), "jar$")[1],
+    package_version = utils::packageVersion("metis.lite")
+  )
+)
 
-    r <- dbSendQuery(conn, statement, ...)
 
-    on.exit(.jcall(r@stat, "V", "close"))
-
-    #message("dbGetQuery()")
+#' Fetch records from a previously executed query
+#'
+#' Fetch the next `n` elements (rows) from the result set and return them
+#' as a data.frame.
+#'
+#' @param res An object inheriting from [DBIResult-class], created by
+#'   [dbSendQuery()].
+#' @param n maximum number of records to retrieve per fetch. Use `n = -1`
+#'   or `n = Inf`
+#'   to retrieve all pending records.  Some implementations may recognize other
+#'   special values.
+#' @param ... Other arguments passed on to methods.
+#' @export
+setMethod(
+  "fetch",
+  signature(res="AthenaResult", n="numeric"),
+  def = function(res, n, block = 1000L, ...) {
 
     nms <- c()
     athena_type_convert <- list()
 
-    cols <- .jcall(r@md, "I", "getColumnCount")
+    cols <- .jcall(res@md, "I", "getColumnCount")
 
     for (i in 1:cols) {
-      ct <- as.character(.jcall(r@md, "I", "getColumnType", i))
+      ct <- as.character(.jcall(res@md, "I", "getColumnType", i))
       athena_type_convert[[i]] <- .jdbc_converters[[ct]]
-      nms <- c(nms, .jcall(r@md, "S", "getColumnLabel", i))
+      nms <- c(nms, .jcall(res@md, "S", "getColumnLabel", i))
     }
 
     athena_type_convert <- set_names(athena_type_convert, nms)
 
-    res <- fetch(r, -1, block = 1000)
+    out <- callNextMethod(res = res, n = n, block = block, ...)
 
     for (nm in names(athena_type_convert)) {
-      res[[nm]] <- athena_type_convert[[nm]](res[[nm]])
+      out[[nm]] <- athena_type_convert[[nm]](out[[nm]])
     }
 
+    out
+
+  }
+)
+
+#' AthenaJDBC
+#'
+#' @param conn Athena connection
+#' @param statement SQL statement
+#' @param ... unused
+#' @importFrom rJava .jcall
+#' @export
+setMethod(
+
+  "dbGetQuery",
+  signature(conn="AthenaConnection", statement="character"),
+
+  definition = function(conn, statement, type_convert=FALSE, ...) {
+
+    r <- dbSendQuery(conn, statement, ...)
+
+    on.exit(.jcall(r@stat, "V", "close"))
+
+    res <- fetch(r, -1, block = 1000L)
+
     class(res) <- c("tbl_df", "tbl", "data.frame")
 
     res
diff --git a/man/fetch-AthenaResult-numeric-method.Rd b/man/fetch-AthenaResult-numeric-method.Rd
new file mode 100644
index 0000000..cde0b0d
--- /dev/null
+++ b/man/fetch-AthenaResult-numeric-method.Rd
@@ -0,0 +1,24 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/z-dbGetQuery.R
+\docType{methods}
+\name{fetch,AthenaResult,numeric-method}
+\alias{fetch,AthenaResult,numeric-method}
+\title{Fetch records from a previously executed query}
+\usage{
+\S4method{fetch}{AthenaResult,numeric}(res, n, block = 1000L, ...)
+}
+\arguments{
+\item{res}{An object inheriting from [DBIResult-class], created by
+[dbSendQuery()].}
+
+\item{n}{maximum number of records to retrieve per fetch. Use `n = -1`
+or `n = Inf`
+to retrieve all pending records.  Some implementations may recognize other
+special values.}
+
+\item{...}{Other arguments passed on to methods.}
+}
+\description{
+Fetch the next `n` elements (rows) from the result set and return them
+as a data.frame.
+}