From 7fee6da44dd506bdfd03aaf672ec90c885fb8520 Mon Sep 17 00:00:00 2001
From: Ben Fasoli <benfasoli@gmail.com>
Date: Wed, 25 Mar 2020 20:29:58 -0700
Subject: [PATCH] Adds support for `JSONB` in PostgresQL Passes dialect to
 instances of `Types` for vendor-specific type mappings.

---
 dataset/database.py |  2 +-
 dataset/types.py    | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/dataset/database.py b/dataset/database.py
index 3155b5a..87d85be 100644
--- a/dataset/database.py
+++ b/dataset/database.py
@@ -45,9 +45,9 @@ class Database(object):
                 if len(schema_qs):
                     schema = schema_qs.pop()
 
-        self.types = Types()
         self.schema = schema
         self.engine = create_engine(url, **engine_kwargs)
+        self.types = Types(self.engine.dialect.name)
         self.url = url
         self.row_type = row_type
         self.ensure_schema = ensure_schema
diff --git a/dataset/types.py b/dataset/types.py
index 97c5b18..904ffa9 100644
--- a/dataset/types.py
+++ b/dataset/types.py
@@ -15,14 +15,26 @@ class Types(object):
     boolean = Boolean
     date = Date
     datetime = DateTime
-    json = JSON
 
-    def guess(self, sample):
+    def __init__(self, dialect = None):
+        self._dialect = dialect
+
+    @property
+    def json(self):
+        if self._dialect is not None and self._dialect == 'postgresql':
+            from sqlalchemy.dialects.postgresql import JSONB
+            return JSONB
+        return JSON
+
+    def guess(self, sample, dialect = None):
         """Given a single sample, guess the column type for the field.
 
         If the sample is an instance of an SQLAlchemy type, the type will be
         used instead.
         """
+        if dialect is not None:
+            self._dialect = dialect
+
         if isinstance(sample, TypeEngine):
             return sample
         if isinstance(sample, bool):