R/db-get_primary_key.R

In tjpalanca/dbtools: Database Tools for R

#' @title Fetch the primary key columns
#'
#' @description We need the primary keys of a database table to provide a
#' sensible default key for upserts and inserts.
#'
#' @param conn (DBIConnection) object
#' @param table (Id) identifier for the primary key
#'
#' @family Extra DBI Methods
#' @export
dbGetPrimaryKey <- function(conn, table) {

  UseMethod("dbGetPrimaryKey")

}

#' @export
dbGetPrimaryKey.PqConnection <- function(conn, table) {

  assert_class(table, "Id")

  if ("schema" %in% names(table@name)) {

    dbGetQuery(
      conn,
      glue_sql(
        "
        SELECT pg_attribute.attname
        FROM pg_index, pg_class, pg_attribute, pg_namespace
        WHERE
        pg_class.oid = {table@name[['table']]}::regclass AND
        indrelid = pg_class.oid AND
        nspname = {table@name[['schema']]} AND
        pg_class.relnamespace = pg_namespace.oid AND
        pg_attribute.attrelid = pg_class.oid AND
        pg_attribute.attnum = any(pg_index.indkey)
        AND indisprimary
        ",
        .con = conn
      )
    ) %>%
      pull(attname)

  } else {

    dbGetQuery(
      conn,
      glue_sql(
        "
        SELECT a.attname
        FROM   pg_index i
        JOIN   pg_attribute a ON a.attrelid = i.indrelid
        AND a.attnum = ANY(i.indkey)
        WHERE  i.indrelid = {table@name[['table']]}::regclass
        AND    i.indisprimary;
        ",
        .con = conn
      )
    ) %>%
      pull(attname)

  }


}

#' @export
dbGetPrimaryKey.SQLiteConnection <- function(conn, table) {

  assert_class(table, "Id")

  dbGetQuery(
    conn,
    glue_sql(
      "PRAGMA table_info({table@name[['table']]})",
      .con = conn
    )
  ) %>%
    filter(pk == 1) %>%
    pull(name)

}