Merge pull request #532 from cody-scott/seed-logic-fix

readded batch logic and included a test

Author: cody-scott

dbt/include/sqlserver/macros/relations/seeds/helpers.sql

@@ -0,0 +1,57 @@
+{% macro sqlserver__get_binding_char() %}
+  {{ return('?') }}
+{% endmacro %}
+
+{% macro sqlserver__get_batch_size() %}
+  {{ return(400) }}
+{% endmacro %}
+
+{% macro calc_batch_size(num_columns) %}
+    {#
+        SQL Server allows for a max of 2098 parameters in a single statement.
+        Check if the max_batch_size fits with the number of columns, otherwise
+        reduce the batch size so it fits.
+    #}
+    {% set max_batch_size = get_batch_size() %}
+    {% set calculated_batch = (2098 / num_columns)|int %}
+    {% set batch_size = [max_batch_size, calculated_batch] | min %}
+
+    {{ return(batch_size) }}
+{%  endmacro %}
+
+{% macro sqlserver__load_csv_rows(model, agate_table) %}
+  {% set cols_sql = get_seed_column_quoted_csv(model, agate_table.column_names) %}
+  {% set batch_size = calc_batch_size(agate_table.column_names|length) %}
+  {% set bindings = [] %}
+  {% set statements = [] %}
+
+  {{ log("Inserting batches of " ~ batch_size ~ " records") }}
+
+  {% for chunk in agate_table.rows | batch(batch_size) %}
+      {% set bindings = [] %}
+
+      {% for row in chunk %}
+          {% do bindings.extend(row) %}
+      {% endfor %}
+
+      {% set sql %}
+          insert into {{ this.render() }} ({{ cols_sql }}) values
+          {% for row in chunk -%}
+              ({%- for column in agate_table.column_names -%}
+                  {{ get_binding_char() }}
+                  {%- if not loop.last%},{%- endif %}
+              {%- endfor -%})
+              {%- if not loop.last%},{%- endif %}
+          {%- endfor %}
+      {% endset %}
+
+      {% do adapter.add_query(sql, bindings=bindings, abridge_sql_log=True) %}
+
+      {% if loop.index0 == 0 %}
+          {% do statements.append(sql) %}
+      {% endif %}
+  {% endfor %}
+
+  {# Return SQL so we can render it out into the compiled files #}
+  {{ return(statements[0]) }}
+{% endmacro %}

tests/functional/adapter/mssql/test_mssql_seed.py

@@ -0,0 +1,41 @@
+import pytest
+from dbt.tests.util import run_dbt
+
+seed_schema_yml = """
+version: 2
+seeds:
+  - name: raw_data
+"""
+
+
+class TestLargeSeed:
+    def build_large_seed_file(self):
+        row_count = 3000
+        column_count = 10
+
+        headers = ",".join(["id"] + [f"column_{_}" for _ in range(1, column_count)])
+        seed_data = [headers]
+        for row in range(1, row_count):
+            row_data = [str(row)]
+            for column in range(1, column_count):
+                row_data += [str(column)]
+
+            row_data = ",".join(row_data)
+            seed_data += [row_data]
+
+        large_seed_file = "\n".join(seed_data)
+        return large_seed_file
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {"name": "generic_tests"}
+
+    @pytest.fixture(scope="class")
+    def seeds(self):
+        return {
+            "raw_data.csv": self.build_large_seed_file(),
+            "schema.yml": seed_schema_yml,
+        }
+
+    def test_large_seed(self, project):
+        run_dbt(["seed"])