Merge pull request #18 from Azure/export-1.1.0

alexkyllo · web-flow · commit d8cfb728fc5f · 2022-12-20T14:50:52.000-08:00
Add export(), slice_sample(), $ for dynamic field access, default auth_type to authorization_code
diff --git a/.Rbuildignore b/.Rbuildignore
@@ -12,3 +12,4 @@ drat.sh
 cran-comments.md
 SECURITY.md
 ^CRAN-SUBMISSION$
+^\.vscode$
diff --git a/.github/workflows/check-standard.yaml b/.github/workflows/check-standard.yaml
@@ -39,11 +39,11 @@ jobs:
             xargs -L1 git config --unset-all
           git push --prune https://token:$token@github.com/${CLOUDYR_REPO}.git +refs/remotes/origin/*:refs/heads/* +refs/tags/*:refs/tags/*
 
-      - uses: r-lib/actions/setup-r@master
+      - uses: r-lib/actions/setup-r@v2
         with:
           r-version: ${{ matrix.config.r }}
 
-      - uses: r-lib/actions/setup-pandoc@master
+      - uses: r-lib/actions/setup-pandoc@v2
 
       - name: Query dependencies
         run: |
diff --git a/.gitignore b/.gitignore
@@ -315,3 +315,5 @@ flycheck_*.el
 .env
 test-creds.*
 inst/doc
+.Renviron
+.vscode/**
diff --git a/CRAN-SUBMISSION b/CRAN-SUBMISSION
@@ -1,3 +1,3 @@
-Version: 1.0.7
-Date: 2022-08-26 15:26:38 UTC
-SHA: f909d6d216d4b10be3fc955132522a4587d8197c
+Version: 1.1.0
+Date: 2022-12-20 21:59:53 UTC
+SHA: 8527843af417cafd880915b0b1ee63e52bb3ff0d
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: AzureKusto
 Title: Interface to 'Kusto'/'Azure Data Explorer'
-Version: 1.0.7
+Version: 1.1.0
 Authors@R: c(
     person("Hong", "Ooi", , "hongooi73@gmail.com", role = "aut"),
     person("Alex", "Kyllo", , "jekyllo@microsoft.com", role = c("aut", "cre")),
diff --git a/NAMESPACE b/NAMESPACE
@@ -19,6 +19,8 @@ S3method(escape,integer64)
 S3method(escape,kql)
 S3method(escape,list)
 S3method(escape,logical)
+S3method(export,kusto_database_endpoint)
+S3method(export,tbl_kusto)
 S3method(filter,tbl_kusto_abstract)
 S3method(full_join,tbl_kusto_abstract)
 S3method(group_by,tbl_kusto_abstract)
@@ -38,6 +40,7 @@ S3method(kql_build,op_mutate)
 S3method(kql_build,op_rename)
 S3method(kql_build,op_select)
 S3method(kql_build,op_set_op)
+S3method(kql_build,op_slice_sample)
 S3method(kql_build,op_summarise)
 S3method(kql_build,op_ungroup)
 S3method(kql_build,op_unnest)
@@ -80,6 +83,7 @@ S3method(semi_join,tbl_kusto_abstract)
 S3method(setdiff,tbl_kusto_abstract)
 S3method(setequal,tbl_kusto_abstract)
 S3method(show_query,tbl_kusto_abstract)
+S3method(slice_sample,tbl_kusto_abstract)
 S3method(summarise,tbl_kusto_abstract)
 S3method(tbl_vars,tbl_kusto_abstract)
 S3method(ungroup,tbl_kusto_abstract)
@@ -97,6 +101,7 @@ export(base_window)
 export(build_kql)
 export(delete_kusto_token)
 export(escape)
+export(export)
 export(filter)
 export(flatten_query)
 export(get_kusto_token)
diff --git a/NEWS.md b/NEWS.md
@@ -1,7 +1,19 @@
+# AzureKusto 1.1.0
+
+* New function `export()` to export query results to Azure Storage.
+* `server` argument of `kusto_database_endpoint` can be just cluster name and no
+  longer needs to be the fully qualified URI. E.g. can pass either
+  `server = "help"` or `server = "https://help.kusto.windows.net"`
+* Changed default `get_kusto_token()` auth_type from "device_code" to
+  "authorization_code" for easier to use sign-in flow.
+* Added `$` as an infix operator in a KQL expression now translates to `.` to
+  enable nested dynamic field access.
+* Added `slice_sample` dplyr verb
+
 # AzureKusto 1.0.7
 
 * Re-release to resolve "Version contains large components" note.
-* 
+
 # AzureKusto 1.0.6.9001
 
 * Regenerate .Rd files for R 4.2+ using updated Roxygen to fix HTML5 issues.
diff --git a/R/endpoint.R b/R/endpoint.R
@@ -93,6 +93,13 @@ kusto_database_endpoint <- function(..., .connection_string=NULL, .query_token=N
     # fix all property names to a given (sub)set, remove quotes from quoted values
     props <- normalize_connstring_properties(props)
 
+    # Make bare cluster name into FQDN for server if it's not already
+    if (!startsWith(props$server, "https://"))
+    {
+        props$server <- paste0("https://", props$server)
+        if (!endsWith(props$server, ".kusto.windows.net"))
+            props$server <- paste0(props$server, ".kusto.windows.net")
+    }
     props$token <- find_endpoint_token(props, .query_token)
 
     props$use_integer64 <- .use_integer64
@@ -109,7 +116,7 @@ normalize_connstring_properties <- function(properties)
     property_list <- list(
         # general properties
         traceclientversion="traceclientversion",
-        server=c("server", "addr", "address", "network address", "datasource", "host"),
+        server=c("server", "addr", "address", "network address", "datasource", "host", "cluster"),
         database=c("database", "initialcatalog", "dbname"),
         tenantid=c("tenantid", "authority"),
         queryconsistency="queryconsistency",
@@ -182,7 +189,7 @@ find_endpoint_token <- function(properties, .query_token)
     {
         message("No app ID supplied; using KustoClient app")
         properties$appclientid <- .kusto_app_id
-        auth_type <- "device_code"
+        auth_type <- "authorization_code"
     }
     else auth_type <- NULL  # KustoClient needs devicecode, otherwise let get_azure_token choose
 
diff --git a/R/kql-build.R b/R/kql-build.R
@@ -197,6 +197,13 @@ kql_build.op_head <- function(op, ...)
     build_kql("take ", kql(escape(n, parens = FALSE)))
 }
 
+#' @export
+kql_build.op_slice_sample <- function(op, ...)
+{
+    n <- lapply(op$args$n, translate_kql)
+    build_kql("sample ", kql(escape(n, parens = FALSE)))
+}
+
 #' @export
 kql_build.op_join <- function(op, ...)
 {
diff --git a/R/kusto_token.R b/R/kusto_token.R
@@ -55,9 +55,10 @@ get_kusto_token <- function(server=NULL, clustername, location=NULL, tenant=NULL
     if(version == 2 && httr::parse_url(server)$path == "")
         server <- c(paste0(sub("/$", "", server), "/.default"), "offline_access", "openid")
 
-    # KustoClient requires devicecode auth if username not supplied
+    # Used to default to "device_code" but "authorization_code" has better ease of use
+    # and also works with .kusto_app_id without supplying a username.
     if(is.null(auth_type) && app == .kusto_app_id && (!"username" %in% names(list(...))))
-        auth_type <- "device_code"
+        auth_type <- "authorization_code"
 
     AzureAuth::get_azure_token(server, tenant, app, auth_type=auth_type, version=version, ...)
 }
diff --git a/R/tbl.R b/R/tbl.R
@@ -84,8 +84,8 @@ group_by.tbl_kusto_abstract <- function(.data, ..., add = FALSE)
     {
         return(.data)
     }
-
-    groups <- group_by_prepare(.data, .dots = dots, add = add)
+    # Updated for dplyr deprecation of .dots and add params
+    groups <- group_by_prepare(.data, !!!dots, .add = add)
     names <- vapply(groups$groups, as_string, character(1))
     add_op_single("group_by",
                   groups$data,
@@ -182,6 +182,12 @@ head.tbl_kusto_abstract <- function(x, n = 6L, ...)
     add_op_single("head", x, args = list(n = n))
 }
 
+#' @export
+slice_sample.tbl_kusto_abstract <- function(x, n = 6L, ...)
+{
+    add_op_single("slice_sample", x, args = list(n = n))
+}
+
 #' Join methods for Kusto tables
 #'
 #' These methods are the same as other joining methods, with the exception of the `.strategy`, `.shufflekeys` and `.num_partitions` optional arguments. They provide hints to the Kusto engine on how to execute the join, and can sometimes be useful to speed up a query. See the Kusto documentation for more details.
@@ -395,6 +401,102 @@ compute.tbl_kusto <- function(tbl, name=generate_table_name(), ...)
     invisible(tbl_kusto(tbl$src, name))
 }
 
+#' Execute the query, store the results in a table, and return a reference to the new table
+#' Run a Kusto query and export results to Azure Storage in Parquet or CSV
+#' format.
+#'
+#' @param query The text of the Kusto query to run
+#' @param storage_uri The URI of the blob storage container to export to
+#' @param name_prefix The filename prefix for each exported file
+#' @param key The account key for the storage container.
+#' uses the identity that is signed into Kusto to authenticate to Azure Storage.
+#' @param format Options are "parquet", "csv", "tsv", "json"
+#' @param distributed logical, indicates whether Kusto should distributed the
+#' export job to multiple nodes, in which case multiple files will be written
+#' to storage concurrently.
+kusto_export_cmd <- function(query, storage_uri, name_prefix, key, format,
+    distributed) {
+    # Make sure the storage uri ends with a slash
+    if (!(format %in% c("parquet", "csv", "tsv", "json")))
+        stop("Format must be one of parquet, csv, tsv, or json.")
+    if (!endsWith(storage_uri, "/"))
+        storage_uri <- paste0(storage_uri, "/")
+    distr <- ifelse(distributed, "true", "false")
+    compr <- ifelse(format == "parquet", "snappy", "gzip")
+    sprintf(".export
+compressed
+to %s (h@'%s%s;%s')
+with (
+sizeLimit=1073741824,
+namePrefix='%s',
+fileExtension='%s',
+compressionType='%s',
+includeHeaders='firstFile',
+encoding='UTF8NoBOM',
+distributed=%s
+)
+<|
+%s
+", format, storage_uri, name_prefix, key, name_prefix, format, compr, distr, query)
+}
+
+#' Execute the Kusto query and export the result to Azure Storage.
+#' @param tbl An object representing a table or database.
+#' @param storage_uri The Azure Storage URI to export files to.
+#' @param query A Kusto query string
+#' @param name_prefix The filename prefix to use for exported files.
+#' @param key default "impersonate" which uses the account signed into Kusto to
+#' authenticate to Azure Storage. An Azure Storage account key.
+#' @param format Options are "parquet", "csv", "tsv", "json"
+#' @param distributed logical, indicates whether Kusto should distributed the
+#' export job to multiple nodes, in which case multiple files will be written
+#' to storage concurrently.
+#' @param ... needed for agreement with generic. Not otherwise used.
+#' @rdname export
+#' @export
+export <- function(tbl, storage_uri, query = NULL, name_prefix = "export",
+    key = "impersonate", format = "parquet", distributed = FALSE, ...) {
+    UseMethod("export")
+}
+
+#' Execute the Kusto query and export the result to Azure Storage.
+#' @param tbl A Kusto database endpoint object, as returned by `kusto_database_endpoint`.
+#' @param query A Kusto query string
+#' @param storage_uri The Azure Storage URI to export files to.
+#' @param name_prefix The filename prefix to use for exported files.
+#' @param key default "impersonate" which uses the account signed into Kusto to
+#' authenticate to Azure Storage. An Azure Storage account key.
+#' @param format Options are "parquet", "csv", "tsv", "json"
+#' @param distributed logical, indicates whether Kusto should distributed the
+#' export job to multiple nodes, in which case multiple files will be written
+#' to storage concurrently.
+#' @param ... needed for agreement with generic. Not otherwise used.
+#' @rdname export
+#' @export
+export.kusto_database_endpoint <- function(tbl, storage_uri, query = NULL, name_prefix = "export",
+    key = "impersonate", format = "parquet", distributed = FALSE, ...) {
+    if (missing(query)) stop("query parameter is required.")
+    is_cmd <- substr(query, 1, 1) == "."
+    if (is_cmd) stop("Management commands cannot be used with export()")
+    q_str <- kusto_export_cmd(query = query, storage_uri = storage_uri,
+        name_prefix = name_prefix, key = key, format = format,
+        distributed = distributed)
+    run_query(tbl, q_str, ...)
+}
+
+#' @rdname export
+#' @export
+export.tbl_kusto <- function(tbl, storage_uri, query = NULL, name_prefix = "export",
+    key = "impersonate", format = "parquet", distributed = FALSE, ...) {
+    database <- tbl$src
+    q <- kql_render(kql_build(tbl))
+    q_str <- kusto_export_cmd(query = q, storage_uri = storage_uri,
+        name_prefix = name_prefix, key = key, format = format,
+        distributed = distributed)
+    res <- run_query(database, q_str, ...)
+    tibble::as_tibble(res)
+}
+
 #' @keywords internal
 #' @export
 print.tbl_kusto_abstract <- function(x, ...)
diff --git a/R/translate-kql.R b/R/translate-kql.R
@@ -319,6 +319,7 @@ base_scalar <- kql_translator(
     `{` = function(x) {
         build_kql("(", x, ")")
     },
+    `$` = kql_infix("."),
     asc = function(x) {
         build_kql(x, kql(" asc"))
     },
diff --git a/README.md b/README.md
diff --git a/man/base_scalar.Rd b/man/base_scalar.Rd
diff --git a/man/export.Rd b/man/export.Rd
diff --git a/man/kusto_export_cmd.Rd b/man/kusto_export_cmd.Rd
diff --git a/tests/testthat/test_translate.r b/tests/testthat/test_translate.r

Original file line number	Diff line number	Diff line change
`@@ -197,6 +197,13 @@ kql_build.op_head <- function(op, ...)`
`197`	`197`	`build_kql("take ", kql(escape(n, parens = FALSE)))`
`198`	`198`	`}`
`199`	`199`
	`200`	`+#' @export`
	`201`	`+kql_build.op_slice_sample <- function(op, ...)`
	`202`	`+{`
	`203`	`+ n <- lapply(op$args$n, translate_kql)`
	`204`	`+ build_kql("sample ", kql(escape(n, parens = FALSE)))`
	`205`	`+}`
	`206`	`+`
`200`	`207`	`#' @export`
`201`	`208`	`kql_build.op_join <- function(op, ...)`
`202`	`209`	`{`