Fix for multiple primary keys

.github/workflows/publish.yaml

@@ -0,0 +1,30 @@
+name: Upload Python Package to PyPI when a Release is Created
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  pypi-publish:
+    name: Publish release to PyPI
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/weiser-ai
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.x"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install setuptools wheel
+      - name: Build package
+        run: |
+          python setup.py sdist bdist_wheel  # Could also be python -m build
+      - name: Publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cube_dbt"
-version = "0.6.1"
+version = "0.6.2"
 description = "dbt integration for Cube"
 authors = [
     {name = "Artyom Keydunov", email = "artyom@cube.dev"},

src/cube_dbt/model.py

@@ -19,16 +19,10 @@ def _init_columns(self) -> None:
       self._detect_primary_key()
 
   def _detect_primary_key(self) -> None:
-    candidates = list(
+    self._primary_key = list(
       column for column in self._columns
       if column.primary_key
     )
-
-    if len(candidates) > 1:
-      column_names = list(column.name for column in candidates)
-      raise RuntimeError(f"More than one primary key column found in {self.name}: {', '.join(column_names)}")
-
-    self._primary_key = candidates[0] if len(candidates) == 1 else None
 
   @property
   def name(self) -> str:
@@ -58,7 +52,7 @@ def column(self, name: str) -> Column:
     return next(column for column in self._columns if column.name == name)
 
   @property
-  def primary_key(self) -> Column or None:
+  def primary_key(self) -> list[Column]:
     self._init_columns()
     return self._primary_key
 

tests/test_model.py

@@ -62,7 +62,7 @@ def test_detect_primary_key_no_columns(self):
       }
     }
     model = Model(model_dict)
-    assert model.primary_key == None
+    assert len(model.primary_key) == 0
 
   def test_detect_primary_key_one_column(self):
     """
@@ -87,8 +87,8 @@ def test_detect_primary_key_one_column(self):
       }
     }
     model = Model(model_dict)
-    assert model.primary_key != None
-    assert model.primary_key.name == 'id'
+    assert len(model.primary_key) == 1
+    assert model.primary_key[0].name == "id"
 
   def test_detect_primary_key_two_columns(self):
     """
@@ -115,8 +115,9 @@ def test_detect_primary_key_two_columns(self):
       }
     }
     model = Model(model_dict)
-    with raises(RuntimeError):
-      model.primary_key
+    assert len(model.primary_key) == 2
+    assert model.primary_key[0].name is not None
+    assert model.primary_key[1].name is not None
 
   def test_as_cube(self):
     model_dict = {
@@ -285,4 +286,37 @@ def test_as_dimensions_render_when_empty(self):
     }
 
     model = Model(model_dict)
-    assert model.as_dimensions() == ''
+    assert model.as_dimensions() == ''
+
+  def test_as_dimensions_render_two_primary_keys(self):
+    model_dict = {
+        "name": "model",
+        "columns": {
+            "id": {
+                "name": "id",
+                "description": "",
+                "meta": {},
+                "data_type": "numeric",
+                "tags": ["primary_key"],
+            },
+            "status": {
+                "name": "status",
+                "description": "",
+                "meta": {},
+                "data_type": None,
+                "tags": ["primary_key"],
+            },
+        },
+    }
+    model = Model(model_dict)
+    assert (
+            model.as_dimensions() == """- name: id
+        sql: id
+        type: number
+        primary_key: true
+      - name: status
+        sql: status
+        type: string
+        primary_key: true
+      """
+    )