feat: add Dataset

Author: Charles Larivier

metabase/__init__.py

@@ -5,3 +5,4 @@
 from metabase.resources.table import Table
 from metabase.resources.segment import Segment
 from metabase.resources.metric import Metric
+from metabase.resources.dataset import Dataset

metabase/resources/dataset.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, List
+
+import pandas as pd
+
+from metabase.resource import CreateResource, Resource
+
+
+class Data(Resource):
+    ENDPOINT = None
+    PRIMARY_KEY = None
+
+    rows: List[List[Any]]
+    cols: dict
+    native_form: dict
+
+    def to_pandas(self):
+        """Returns the query results as a Pandas DataFrame."""
+        columns = [col["display_name"] for col in self.cols]
+        return pd.DataFrame(data=self.rows, columns=columns)
+
+
+class Dataset(CreateResource):
+    ENDPOINT = "/api/dataset"
+    PRIMARY_KEY = None
+
+    context: str
+    status: str
+    database_id: int
+    data: Data
+    row_count: int
+
+    started_at: str
+    running_time: int
+    json_query: str
+    average_execution_time: int = None
+
+    @classmethod
+    def create(
+            cls,
+            database: int,
+            type: str,
+            query: dict,
+            **kwargs
+    ) -> Dataset:
+        dataset = super(Dataset, cls).create(database=database, type=type, query=query)
+        dataset.data = Data(**dataset.data)
+        return dataset
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Returns the query results as a Pandas DataFrame."""
+        return self.data.to_pandas()

tests/resources/test_dataset.py

@@ -0,0 +1,41 @@
+from metabase.resources.dataset import Data, Dataset
+
+from tests.helpers import IntegrationTestCase
+
+import pandas as pd
+
+
+class DatasetTests(IntegrationTestCase):
+    def test_import(self):
+        """Ensure Metric can be imported from Metabase."""
+        from metabase import Dataset
+        self.assertIsNotNone(Dataset())
+
+    def test_create(self):
+        """Ensure Dataset.create() executes a query and returns a Dataset instance with the query results."""
+        dataset = Dataset.create(
+            database=1,
+            type="query",
+            query={
+                "source-table": 1,
+                "breakout": [["field", 7, {"temporal-unit": "year"}]],
+                "aggregation": [["count"]]
+            }
+        )
+        self.assertIsInstance(dataset, Dataset)
+        self.assertIsInstance(dataset.data, Data)
+
+    def test_to_pandas(self):
+        """Ensure Dataset.to_pandas() returns a Pandas DataFrame."""
+        dataset = Dataset.create(
+            database=1,
+            type="query",
+            query={
+                "source-table": 1,
+                "breakout": [["field", 7, {"temporal-unit": "year"}]],
+                "aggregation": [["count"]]
+            }
+        )
+        df = dataset.to_pandas()
+        self.assertIsInstance(df, pd.DataFrame)
+        self.assertListEqual(["Created At", "Count"], df.columns.tolist())