Feature/python model v1

dbt/adapters/snowflake/__version__.py

@@ -1 +1 @@
-version = "1.2.0a1"
+version = "1.3.0a1"

dbt/adapters/snowflake/impl.py

@@ -4,6 +4,7 @@
 import agate
 
 from dbt.adapters.base.impl import AdapterConfig
+from dbt.adapters.base.meta import available
 from dbt.adapters.sql import SQLAdapter  # type: ignore
 from dbt.adapters.sql.impl import (
     LIST_SCHEMAS_MACRO_NAME,
@@ -160,3 +161,33 @@ def quote_seed_column(self, column: str, quote_config: Optional[bool]) -> str:
 
     def timestamp_add_sql(self, add_to: str, number: int = 1, interval: str = "hour") -> str:
         return f"DATEADD({interval}, {number}, {add_to})"
+
+    @available.parse_none
+    def submit_python_job(self, parsed_model:dict, compiled_code: str):
+        schema = getattr(parsed_model, "schema", self.config.credentials.schema)
+        database = getattr(parsed_model, "database", self.config.credentials.database)
+        identifier = parsed_model['alias']
+        proc_name = f"{database}.{schema}.{identifier}__dbt_sp"
+        packages = ['snowflake-snowpark-python'] + parsed_model['config'].get('packages', [])
+        packages = "', ".join(packages)
+        python_stored_procedure = f"""
+CREATE OR REPLACE PROCEDURE {proc_name} ()
+RETURNS STRING
+LANGUAGE PYTHON
+RUNTIME_VERSION = '3.8' -- TODO should this be configurable?
+PACKAGES = ('{packages}')
+HANDLER = 'main'
+AS
+$$
+{compiled_code}
+
+$$;
+        """
+        self.execute(python_stored_procedure, auto_begin=False, fetch=False)
+        self.execute(f"CALL {proc_name}()", auto_begin=False, fetch=False)
+        self.execute(f"drop procedure if exists {proc_name}(string)", auto_begin=False, fetch=False)
+
+        # TODO add proper return
+        return "OK"
+
+

dbt/include/snowflake/macros/adapters.sql

@@ -1,41 +1,49 @@
-{% macro snowflake__create_table_as(temporary, relation, sql) -%}
-  {%- set transient = config.get('transient', default=true) -%}
-  {%- set cluster_by_keys = config.get('cluster_by', default=none) -%}
-  {%- set enable_automatic_clustering = config.get('automatic_clustering', default=false) -%}
-  {%- set copy_grants = config.get('copy_grants', default=false) -%}
-
-  {%- if cluster_by_keys is not none and cluster_by_keys is string -%}
-    {%- set cluster_by_keys = [cluster_by_keys] -%}
-  {%- endif -%}
-  {%- if cluster_by_keys is not none -%}
-    {%- set cluster_by_string = cluster_by_keys|join(", ")-%}
-  {% else %}
-    {%- set cluster_by_string = none -%}
-  {%- endif -%}
-  {%- set sql_header = config.get('sql_header', none) -%}
-
-  {{ sql_header if sql_header is not none }}
-
-      create or replace {% if temporary -%}
-        temporary
-      {%- elif transient -%}
-        transient
-      {%- endif %} table {{ relation }} {% if copy_grants and not temporary -%} copy grants {%- endif %} as
-      (
-        {%- if cluster_by_string is not none -%}
-          select * from(
-            {{ sql }}
-            ) order by ({{ cluster_by_string }})
-        {%- else -%}
-          {{ sql }}
-        {%- endif %}
-      );
-    {% if cluster_by_string is not none and not temporary -%}
-      alter table {{relation}} cluster by ({{cluster_by_string}});
+{% macro snowflake__create_table_as(temporary, relation, compiled_code, language='sql') -%}
+  {%- if language == 'sql' -%}
+    {%- set transient = config.get('transient', default=true) -%}
+    {%- set cluster_by_keys = config.get('cluster_by', default=none) -%}
+    {%- set enable_automatic_clustering = config.get('automatic_clustering', default=false) -%}
+    {%- set copy_grants = config.get('copy_grants', default=false) -%}
+
+    {%- if cluster_by_keys is not none and cluster_by_keys is string -%}
+      {%- set cluster_by_keys = [cluster_by_keys] -%}
     {%- endif -%}
-    {% if enable_automatic_clustering and cluster_by_string is not none and not temporary  -%}
-      alter table {{relation}} resume recluster;
+    {%- if cluster_by_keys is not none -%}
+      {%- set cluster_by_string = cluster_by_keys|join(", ")-%}
+    {% else %}
+      {%- set cluster_by_string = none -%}
     {%- endif -%}
+    {%- set sql_header = config.get('sql_header', none) -%}
+
+    {{ sql_header if sql_header is not none }}
+
+        create or replace {% if temporary -%}
+          temporary
+        {%- elif transient -%}
+          transient
+        {%- endif %} table {{ relation }} {% if copy_grants and not temporary -%} copy grants {%- endif %} as
+        (
+          {%- if cluster_by_string is not none -%}
+            select * from(
+              {{ compiled_code }}
+              ) order by ({{ cluster_by_string }})
+          {%- else -%}
+            {{ compiled_code }}
+          {%- endif %}
+        );
+      {% if cluster_by_string is not none and not temporary -%}
+        alter table {{relation}} cluster by ({{cluster_by_string}});
+      {%- endif -%}
+      {% if enable_automatic_clustering and cluster_by_string is not none and not temporary  -%}
+        alter table {{relation}} resume recluster;
+      {%- endif -%}
+
+  {%- elif language == 'python' -%}
+    {{ py_complete_script(compiled_code=compiled_code, target_relation=relation) }}
+  {%- else -%}
+      {% do exceptions.raise_compiler_error("snowflake__create_table_as macro didn't get supported language") %}
+  {%- endif -%}
+
 {% endmacro %}
 
 {% macro get_column_comment_sql(column_name, column_dict) %}

dbt/include/snowflake/macros/materializations/incremental.sql

@@ -28,9 +28,10 @@
 
   {% set original_query_tag = set_query_tag() %}
 
+  {#-- Set vars --#}
   {%- set unique_key = config.get('unique_key') -%}
   {%- set full_refresh_mode = (should_full_refresh()) -%}
-
+  {%- set language = config.get('language') -%}
   {% set target_relation = this %}
   {% set existing_relation = load_relation(this) %}
   {% set tmp_relation = make_temp_relation(this) %}
@@ -42,19 +43,27 @@
   {{ run_hooks(pre_hooks) }}
 
   {% if existing_relation is none %}
-    {% set build_sql = create_table_as(False, target_relation, sql) %}
+    {%- call statement('main', language=language) -%}
+      {{ create_table_as(False, target_relation, compiled_code, language) }}
+    {%- endcall -%}
 
   {% elif existing_relation.is_view %}
     {#-- Can't overwrite a view with a table - we must drop --#}
     {{ log("Dropping relation " ~ target_relation ~ " because it is a view and this model is a table.") }}
     {% do adapter.drop_relation(existing_relation) %}
-    {% set build_sql = create_table_as(False, target_relation, sql) %}
-
+    {%- call statement('main', language=language) -%}
+      {{ create_table_as(False, target_relation, compiled_code, language) }}
+    {%- endcall -%}
   {% elif full_refresh_mode %}
-    {% set build_sql = create_table_as(False, target_relation, sql) %}
+    {%- call statement('main', language=language) -%}
+      {{ create_table_as(False, target_relation, compiled_code, language) }}
+    {%- endcall -%}
 
   {% else %}
-    {% do run_query(create_table_as(True, tmp_relation, sql)) %}
+    {%- call statement('create_tmp_relation', language=language) -%}
+      {{ create_table_as(True, tmp_relation, compiled_code, language) }}
+    {%- endcall -%}
+
     {% do adapter.expand_target_column_types(
            from_relation=tmp_relation,
            to_relation=target_relation) %}
@@ -63,7 +72,9 @@
     {% if not dest_columns %}
       {% set dest_columns = adapter.get_columns_in_relation(existing_relation) %}
     {% endif %}
-    {% set build_sql = dbt_snowflake_get_incremental_sql(strategy, tmp_relation, target_relation, unique_key, dest_columns) %}
+     {%- call statement('main') -%}
+      {{ dbt_snowflake_get_incremental_sql(strategy, tmp_relation, target_relation, unique_key, dest_columns) }}
+    {%- endcall -%}
 
   {% endif %}
 

dbt/include/snowflake/macros/materializations/table.sql

@@ -3,6 +3,7 @@
   {% set original_query_tag = set_query_tag() %}
 
   {%- set identifier = model['alias'] -%}
+  {%- set language = config.get('language') -%}
 
   {%- set old_relation = adapter.get_relation(database=database, schema=schema, identifier=identifier) -%}
   {%- set target_relation = api.Relation.create(identifier=identifier,
@@ -18,9 +19,8 @@
     {{ drop_relation_if_exists(old_relation) }}
   {% endif %}
 
-  --build model
-  {% call statement('main') -%}
-    {{ create_table_as(false, target_relation, sql) }}
+  {% call statement('main', language=language) -%}
+      {{ create_table_as(False, target_relation, compiled_code, language) }}
   {%- endcall %}
 
   {{ run_hooks(post_hooks) }}
@@ -32,3 +32,24 @@
   {{ return({'relations': [target_relation]}) }}
 
 {% endmaterialization %}
+
+{% macro py_complete_script(compiled_code, target_relation) %}
+{{ compiled_code }}
+def materialize(session, df, target_relation):
+    # we have to make sure pandas is imported
+    import pandas
+    if isinstance(df, pandas.core.frame.DataFrame):
+        # session.write_pandas does not have overwrite function
+        df = session.createDataFrame(df)
+    df.write.mode("overwrite").save_as_table("{{ target_relation }}")
+
+def main(session):
+    """
+    TODOs:
+      - what should this return? can we make a real RunResult?
+    """
+    dbt = dbtObj(session.table)
+    df = model(dbt, session)
+    materialize(session, df, dbt.this)
+    return "OK"
+{% endmacro %}

setup.py

@@ -50,7 +50,7 @@ def _get_dbt_core_version():
 
 
 package_name = "dbt-snowflake"
-package_version = "1.2.0a1"
+package_version = "1.3.0a1"
 dbt_core_version = _get_dbt_core_version()
 description = """The Snowflake adapter plugin for dbt"""
 

tests/functional/adapter/test_basic.py

@@ -13,6 +13,7 @@
 from dbt.tests.adapter.basic.test_snapshot_timestamp import BaseSnapshotTimestamp
 from dbt.tests.adapter.basic.test_adapter_methods import BaseAdapterMethod
 
+
 class TestSimpleMaterializationsSnowflake(BaseSimpleMaterializations):
     pass
 
@@ -51,4 +52,4 @@ class TestSnapshotTimestampSnowflake(BaseSnapshotTimestamp):
 class TestBaseAdapterMethodSnowflake(BaseAdapterMethod):
     @pytest.fixture(scope="class")
     def equal_tables(self):
-        return ["MODEL", "EXPECTED"]
+        return ["MODEL", "EXPECTED"]

tests/functional/adapter/test_python_model.py

@@ -0,0 +1,69 @@
+import pytest
+from dbt.tests.util import run_dbt, write_file
+from dbt.tests.adapter.python_model.test_python_model import BasePythonModelTests
+
+basic_sql = """
+select 1 as id union all
+select 1 as id union all
+select 1 as id union all
+select 1 as id union all
+select 1 as id
+"""
+basic_python = """
+def model(dbt, session):
+    dbt.config(
+        materialized='table',
+    )
+    df =  dbt.ref("my_sql_model")
+    df = df.limit(2)
+    return df
+"""
+
+class TestBasePythonModelSnowflake:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "my_sql_model.sql": basic_sql,
+            "my_python_model.py": basic_python,
+        }
+
+    def test_singular_tests(self, project):
+        # test command
+        results = run_dbt(["run"])
+        assert len(results) == 2
+
+
+models__simple_python_model = """
+import pandas
+
+def model(dbt, session):
+    dbt.config(
+        materialized='table',
+    )
+    data = [[1,2]] * 10
+    return pandas.DataFrame(data, columns=['test', 'test2'])
+"""
+models__simple_python_model_v2 = """
+import pandas
+
+def model(dbt, session):
+    dbt.config(
+        materialized='table',
+    )
+    data = [[1,2]] * 10
+    return pandas.DataFrame(data, columns=['test1', 'test3'])
+"""
+
+class TestChangingSchemaSnowflake:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "simple_python_model.py": models__simple_python_model
+            }
+    def test_changing_schema(self,project):
+        run_dbt(["run"])
+        write_file(models__simple_python_model_v2, project.project_root + '/models', "simple_python_model.py")
+        run_dbt(["run"])
+
+
+