From 921940da6174aa23d5580d76fed7dcd2184dd5ae Mon Sep 17 00:00:00 2001
From: DeflateAwning <11021263+DeflateAwning@users.noreply.github.com>
Date: Thu, 13 Feb 2025 20:47:11 -0700
Subject: [PATCH] Add support for reading CSV files (#8)

---
 package.json           |  3 +++
 src/parquetDocument.ts | 29 +++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/package.json b/package.json
index 893736f..96b8a84 100644
--- a/package.json
+++ b/package.json
@@ -65,6 +65,9 @@
                     },
                     {
                         "filenamePattern": "*.pq"
+                    },
+                    {
+                        "filenamePattern": "*.csv"
                     }
                 ],
                 "priority": "default"
diff --git a/src/parquetDocument.ts b/src/parquetDocument.ts
index ba06b05..4ab9589 100644
--- a/src/parquetDocument.ts
+++ b/src/parquetDocument.ts
@@ -1,10 +1,12 @@
 import * as vscode from 'vscode';
 import { Disposable } from './dispose';
 import { getNonce } from './util';
-import { parse } from "path"
-
+import { parse, extname } from "path"
 import * as duckdb from 'duckdb';
 
+const CSV_EXTENSIONS = [".csv"];
+const PARQUET_EXTENSIONS = [".pq", ".parq", ".parquet"];
+
 type IMessage = {
     type: 'query' | 'more';
     success: boolean;
@@ -14,7 +16,7 @@ type IMessage = {
 } | { type: 'config', autoQuery: boolean };
 
 /**
- * Define the document (the data model) used for paw draw files.
+ * Define the document (the data model) used for table files.
  */
 class ParquetDocument extends Disposable implements vscode.CustomDocument {
 
@@ -35,14 +37,25 @@ class ParquetDocument extends Disposable implements vscode.CustomDocument {
         this._uri = uri;
         this._db = new duckdb.Database(':memory:');
 
-        const config = vscode.workspace.getConfiguration('flat-file-explorer')
+        const config = vscode.workspace.getConfiguration('flat-file-explorer');
         let tableName: string = config.get("tableName")!;
         if (config.get("useFileNameAsTableName"))
-            tableName = parse(uri.fsPath).name
+            tableName = parse(uri.fsPath).name;
+
+        const fileExtension = extname(uri.fsPath).toLowerCase();
+        let query = "";
+
+        if (CSV_EXTENSIONS.includes(fileExtension)) {
+            query = `CREATE VIEW ${tableName} AS SELECT * FROM read_csv('${uri.fsPath}');`;
+        } else if (PARQUET_EXTENSIONS.includes(fileExtension)) {
+            query = `CREATE VIEW ${tableName} AS SELECT * FROM read_parquet('${uri.fsPath}');`;
+        } else {
+            // If this error occurs, check that the trigger types in `package.json` are in sync
+            // with the extension definition arrays at the top of this file.
+            throw new Error("Unsupported file type. Should not have opened with Flat File Explorer.");
+        }
 
-        this.db.exec(
-            `CREATE VIEW ${tableName} AS SELECT * FROM read_parquet('${uri.fsPath}');`
-        );
+        this.db.exec(query);
     }
 
     public get uri() { return this._uri; }