Learn keys from SQLite databases

R/db-helpers.R

@@ -163,6 +163,11 @@ is_postgres <- function(dest) {
     inherits(dest, "PqConnection")
 }
 
+is_sqlite <- function(dest) {
+  inherits(dest, "SQLiteConnection") ||
+    inherits(dest, "src_SQLiteConnection")
+}
+
 src_from_src_or_con <- function(dest) {
   if (is.src(dest)) dest else dbplyr::src_dbi(dest)
 }

R/learn.R

@@ -85,6 +85,9 @@ db_learn_query <- function(dest, ...) {
   if (is_postgres(dest)) {
     return(postgres_learn_query(dest, ...))
   }
+  if (is_sqlite(dest)) {
+    return(sqlite_learn_query())
+  }
 }
 
 mssql_learn_query <- function() { # taken directly from {datamodelr}
@@ -181,6 +184,34 @@ postgres_learn_query <- function(con, schema = "public", table_type = "BASE TABL
   )
 }
 
+sqlite_learn_query <- function() {
+  "SELECT
+    NULL as schema,
+    c.*,
+    fk.ref,
+    fk.ref_col
+  FROM (
+    SELECT
+      m.name as 'table',
+      c.name as column,
+      c.pk as 'key',
+      c.'type',
+      c.'notnull' as mandatory,
+      c.cid + 1 as column_order
+    FROM (SELECT * FROM sqlite_master UNION ALL SELECT * FROM sqlite_temp_master) m
+    JOIN pragma_table_info(m.name) c ON m.type = 'table'
+  ) c
+  LEFT JOIN (
+    SELECT
+      m.name as 'table',
+      fk.'table' as ref,
+      fk.'from' as column,
+      fk.'to' as ref_col
+    FROM (SELECT * FROM sqlite_master UNION ALL SELECT * FROM sqlite_temp_master) m
+    JOIN pragma_foreign_key_list(m.name) fk
+  ) fk  ON c.'table' = fk.'table' AND c.column = fk.column"
+}
+
 # FIXME: only needed for `dm_learn_from_db()` <- needs to be implemented in a different manner
 legacy_new_dm <- function(tables = NULL, data_model = NULL) {
   if (is_null(tables) && is_null(data_model)) {

tests/testthat/test-learn.R

@@ -84,3 +84,43 @@ test_that("Learning from SQLite works (#288)?", {
     dm(test = tibble(a = 1:3))
   )
 })
+
+test_that("Learning keys from SQLite works", {
+  src_sqlite <- skip_if_error(src_sqlite(":memory:", TRUE))
+  con_sqlite <- src_sqlite$con
+  on.exit(dbDisconnect(con_sqlite))
+
+  dbExecute(con_sqlite, "CREATE TABLE first (id INTEGER PRIMARY KEY)")
+  dbExecute(con_sqlite, paste(
+    "CREATE TABLE second (",
+    "id INTEGER PRIMARY KEY, first_id INTEGER,",
+    "FOREIGN KEY (first_id) REFERENCES first (id))"
+  ))
+  dbExecute(con_sqlite, paste(
+    "CREATE TABLE third (",
+    "id INTEGER PRIMARY KEY, first_id INTEGER, second_id INTEGER,",
+    "FOREIGN KEY (first_id) REFERENCES first (id),",
+    "FOREIGN KEY (second_id) REFERENCES second (id))"
+  ))
+
+  x <- dm_from_src(src_sqlite, learn_keys = TRUE)
+  expect_equivalent(
+    dplyr::mutate_all(dm_get_all_fks(x), format, justify = "none"),
+    tibble::tribble(
+      ~child_table, ~child_fk_cols, ~parent_table,
+          "second",     "first_id",       "first",
+           "third",     "first_id",       "first",
+           "third",    "second_id",      "second"
+    )
+  )
+
+  expect_equivalent(
+    dplyr::mutate_all(dm_get_all_pks(x), format, justify = "none"),
+    tibble::tribble(
+        ~table, ~pk_col,
+       "first",    "id",
+      "second",    "id",
+       "third",    "id"
+    )
+  )
+})