Add Snowflake support (#29)

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
     args: ["--line-length", "119"]
     stages: [commit]
 - repo: https://github.com/pycqa/flake8
-  rev: 3.9.2
+  rev: 7.0.0
   hooks:
   - id: flake8
     args: ["--ignore", "E203,E266,E501,W503", "--max-line-length", "119", "--max-complexity", "18", "--select", "B,C,E,F,W,T4,B9"]

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 * Support for Redshift
+* Support for Snowflake
 
 ### Changed
 

README.md

@@ -16,6 +16,7 @@ The library currently supports the following databases.
 * BigQuery
 * Clickhouse
 * Redshift
+* Snowflake
 
 ## Installation
 
@@ -30,6 +31,9 @@ pip install --upgrade "sql-mock[clickhouse]"
 
 # Redshift
 pip install --upgrade "sql-mock[redshift]"
+
+# Snowflake
+pip install --upgrade "sql-mock[snowflake]"
 ```
 
 If you need to modify this source code, install the dependencies using poetry:

docsource/getting_started/installation.md

@@ -15,6 +15,9 @@ pip install --upgrade "sql-mock[clickhouse]"
 
 # Redshift
 pip install --upgrade "sql-mock[redshift]"
+
+# Snowflake
+pip install --upgrade "sql-mock[snowflake]"
 ```
 
 If you need to modify this source code, install the dependencies using poetry:

docsource/index.rst

@@ -43,6 +43,7 @@ It provides a consistent and convenient way to test the execution of your query
    usage/bigquery/index
    usage/clickhouse/index
    usage/redshift/index
+   usage/snowflake/index
 
 .. toctree::
    :maxdepth: 3

docsource/usage/snowflake/examples.md

@@ -0,0 +1,63 @@
+```{toctree}
+:maxdepth: 2
+```
+
+# Example: Testing Subscription Counts in Snowflake
+
+```python
+import datetime
+from sql_mock.snowflake import column_mocks as col
+from sql_mock.snowflake.table_mocks import SnowflakeMockTable
+from sql_mock.table_mocks import table_meta
+
+# Define mock tables for your data model that inherit from SnowflakeMockTable
+@table_meta(table_ref="data.users")
+class UserTable(SnowflakeMockTable):
+    user_id = col.INTEGER(default=1)
+    user_name = col.STRING(default="Mr. T")
+
+
+@table_meta(table_ref="data.subscriptions")
+class SubscriptionTable(SnowflakeMockTable):
+    subscription_id = col.INTEGER(default=1)
+    period_start_date = col.DATE(default=datetime.date(2023, 9, 5))
+    period_end_date = col.DATE(default=datetime.date(2023, 9, 5))
+    user_id = col.INTEGER(default=1)
+
+
+# Define a mock table for your expected results
+class SubscriptionCountTable(SnowflakeMockTable):
+    subscription_count = col.INTEGER(default=1)
+    user_id = col.INTEGER(default=1)
+
+# Your original SQL query
+query = """
+SELECT
+    count(*) AS subscription_count,
+    user_id
+FROM data.users
+LEFT JOIN data.subscriptions USING(user_id)
+GROUP BY user_id
+"""
+
+def test_something():
+  # Create mock data for the 'data.users' and 'data.subscriptions' tables
+  users = UserTable.from_dicts([{'user_id': 1}, {'user_id': 2}])
+  subscriptions = SubscriptionTable.from_dicts([
+      {'subscription_id': 1, 'user_id': 1},
+      {'subscription_id': 2, 'user_id': 1},
+      {'subscription_id': 2, 'user_id': 2},
+  ])
+
+  # Define your expected results
+  expected = [
+      {"USER_ID": 1, "SUBSCRIPTION_COUNT": 2},
+      {"USER_ID": 2, "SUBSCRIPTION_COUNT": 1},
+  ]
+
+  # Simulate the SQL query using SQL Mock
+  res = SubscriptionCountTable.from_mocks(query=query, input_data=[users, subscriptions])
+
+  # Assert the results
+  res.assert_equal(expected)
+```

docsource/usage/snowflake/index.rst

@@ -0,0 +1,10 @@
+Snowflake
+=====================
+
+This section documents the specifics on how to use SQL Mock with Snowflake
+
+.. toctree::
+   :maxdepth: 4
+
+   ./settings.md
+   ./examples.md

docsource/usage/snowflake/settings.md

@@ -0,0 +1,13 @@
+```{toctree}
+:maxdepth: 2
+```
+
+# Settings
+
+In order to use SQL Mock with Snowflake, you need to provide the following environment variables when you run tests:
+
+* `SQL_MOCK_SNOWFLAKE_ACCOUNT`: The name of your Snowflake account
+* `SQL_MOCK_SNOWFLAKE_USER`: The name of your Snowflake user
+* `SQL_MOCK_SNOWFLAKE_PASSWORD`: The password for your Snowflake user
+
+Having those environment variables enables SQL Mock to connect to your Snowflake instance.

examples/snowflake/test_example.py

@@ -0,0 +1,87 @@
+import datetime
+
+from sql_mock.snowflake import column_mocks as col
+from sql_mock.snowflake.table_mocks import SnowflakeMockTable
+from sql_mock.table_mocks import table_meta
+
+
+@table_meta(table_ref="data.users")
+class UserTable(SnowflakeMockTable):
+    user_id = col.INTEGER(default=1)
+    user_name = col.STRING(default="Mr. T")
+
+
+@table_meta(table_ref="data.subscriptions")
+class SubscriptionTable(SnowflakeMockTable):
+    subscription_id = col.INTEGER(default=1)
+    period_start_date = col.DATE(default=datetime.date(2023, 9, 5))
+    period_end_date = col.DATE(default=datetime.date(2023, 9, 5))
+    user_id = col.INTEGER(default=1)
+
+
+@table_meta(
+    query_path="./examples/test_query.sql",
+    default_inputs=[
+        UserTable([]),
+        SubscriptionTable([]),
+    ],  # We can provide defaults for the class if needed. Then we don't always need to provide data for all input tables.
+)
+class MultipleSubscriptionUsersTable(SnowflakeMockTable):
+    user_id = col.INTEGER(default=1)
+
+
+def test_something():
+    users = UserTable.from_dicts([{"user_id": 1}, {"user_id": 2}])
+    subscriptions = SubscriptionTable.from_dicts(
+        [
+            {"subscription_id": 1, "user_id": 1},
+            {"subscription_id": 2, "user_id": 1},
+            {"subscription_id": 2, "user_id": 2},
+        ]
+    )
+
+    subscriptions_per_user__expected = [
+        {"USER_ID": 1, "SUBSCRIPTION_COUNT": 2},
+        {"USER_ID": 2, "SUBSCRIPTION_COUNT": 1},
+    ]
+    users_with_multiple_subs__expected = [{"USER_ID": 1, "SUBSCRIPTION_COUNT": 2}]
+    end_result__expected = [{"USER_ID": 1}]
+
+    res = MultipleSubscriptionUsersTable.from_mocks(input_data=[users, subscriptions])
+
+    # Check the results of the subscriptions_per_user CTE
+    res.assert_cte_equal("subscriptions_per_user", subscriptions_per_user__expected)
+    # Check the results of the users_with_multiple_subs CTE
+    res.assert_cte_equal("users_with_multiple_subs", users_with_multiple_subs__expected)
+    # Check the end result
+    res.assert_equal(end_result__expected)
+
+
+def test_with_defaults_for_subscriptions_table():
+    """
+    In this test case we don't provide a mock for subscriptions
+    because we use the class default Subscriptions([]) which translates to an empty table.
+    """
+    users = UserTable.from_dicts(
+        [
+            {"user_id": 1},
+            {"user_id": 2},
+        ]
+    )
+
+    subscriptions_per_user__expected = [
+        {"USER_ID": 1, "SUBSCRIPTION_COUNT": 0},
+        {"USER_ID": 2, "SUBSCRIPTION_COUNT": 0},
+    ]
+    users_with_multiple_subs__expected = []
+    end_result__expected = []
+
+    # We don't provide a mock input for subscriptions
+    res = MultipleSubscriptionUsersTable.from_mocks(input_data=[users])
+
+    # Check the results of the subscriptions_per_user CTE
+    res.assert_cte_equal("subscriptions_per_user", subscriptions_per_user__expected)
+    # Check the results of the users_with_multiple_subs CTE
+    res.assert_cte_equal("users_with_multiple_subs", users_with_multiple_subs__expected)
+    # Check the end result
+    res.assert_equal(end_result__expected)