DH-4949/add support for duckdb (#261)

* DH-4949/add support for duckdb * DH4949/refromat
Dataherald · Nov 21, 2023 · f3d42df · f3d42df
1 parent 71ec83d
commit f3d42df
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -226,6 +226,12 @@ You can generate the `connection_uri` parameter in the API call for each of the
 "connection_uri": awsathena+rest://<aws_access_key_id>:<aws_secret_access_key>@athena.<region_name>.amazonaws.com:443/<schema_name>?s3_staging_dir=<s3_staging_dir>&work_group=primary
 ```
 
+**DuckDB**
+To connect to DuckDB you should first host your database in MotherDuck and fetch your authentication token.
+```
+"connection_uri": duckdb:///md:<database_name>?motherduck_token=<your_token>
+```
+
 **BigQuery**
 To connect to BigQuery you should create a json credential file. Please follow Steps 1-3 under "Configure 
 BigQuery Authentication in Google Cloud Platform" in 

diff --git a/dataherald/sql_database/base.py b/dataherald/sql_database/base.py
@@ -62,6 +62,9 @@ def from_uri(
     ) -> "SQLDatabase":
         """Construct a SQLAlchemy engine from URI."""
         _engine_args = engine_args or {}
+        if database_uri.lower().startswith("duckdb"):
+            config = {"autoload_known_extensions": False}
+            _engine_args["connect_args"] = {"config": config}
         engine = create_engine(database_uri, **_engine_args)
         return cls(engine, **kwargs)
 

diff --git a/requirements.txt b/requirements.txt
@@ -35,3 +35,5 @@ boto3==1.28.38
 botocore==1.31.38
 PyAthena==3.0.6
 tiktoken==0.5.1
+duckdb-engine==0.9.1
+duckdb==0.9.1