Merge pull request #52 from pageuppeople-opensource/fix/compound-primary-keys

Fixed bug with handling batches of data with compound PKs

Author: ChintanRaval

rdl/data_sources/MsSqlDataSource.py

@@ -231,17 +231,22 @@ def __build_where_clause(batch_key_tracker, table_alias):
 
         try:
             sql_builder = io.StringIO()
+            where_stack = io.StringIO()
+            where_stack.write("1 = 1")
             for primary_key in batch_key_tracker.bookmarks:
                 if has_value:
-                    sql_builder.write(" AND ")
+                    sql_builder.write(" OR")
 
                 sql_builder.write(
-                    f" {table_alias}.{primary_key} > {batch_key_tracker.bookmarks[primary_key]}")
+                    f" ({where_stack.getvalue()} AND {table_alias}.{primary_key} > {batch_key_tracker.bookmarks[primary_key]})")
                 has_value = True
+                where_stack.write(
+                    f" AND {table_alias}.{primary_key} = {batch_key_tracker.bookmarks[primary_key]}")
 
             return sql_builder.getvalue()
         finally:
             sql_builder.close()
+            where_stack.close()
 
     @staticmethod
     def __build_change_table_on_clause(batch_key_tracker):

tests/integration_tests/mssql_source/assertions/compound_pk_test_full_refresh_assertions.sql

@@ -0,0 +1,16 @@
+SET client_encoding
+TO 'UTF8';
+
+DO $$
+BEGIN
+  IF ((SELECT COUNT(*)
+    FROM rdl_integration_tests.load_compound_pk ) = 4)
+    OR
+    ((SELECT COUNT(*)
+    FROM rdl_integration_tests.load_compound_pk ) = 10) THEN
+    RAISE NOTICE '[COMPOUND KEY MSSQL IMPORT TEST] PASS';
+ELSE
+    RAISE EXCEPTION '[COMPOUND KEY MSSQL IMPORT TEST] FAIL: Did not find the required 4 (or 10) rows.';
+END
+IF;
+END $$;

tests/integration_tests/mssql_source/config/CompoundPkTest.json

@@ -9,7 +9,7 @@
   "load_table": "load_compound_pk",
 
   "batch": {
-    "size": 100000
+    "size": 2
   },
   "columns": [
     {

tests/integration_tests/mssql_source/source_database_setup/change_compound_pk.sql

@@ -4,5 +4,13 @@ INSERT CompoundPk
       Id2
       )
       SELECT 3, 5
+UNION ALL
+      SELECT 3, 6
+UNION ALL
+      SELECT 3, 7
+UNION ALL
+      SELECT 4, 3
+UNION ALL
+      SELECT 4, 4
 UNION ALL
       SELECT 7, 8

tests/integration_tests/test_full_refresh_from_mssql.cmd

@@ -7,3 +7,6 @@ if %errorlevel% neq 0 exit /b %errorlevel%
 
 psql -U postgres -d rdl_integration_test_target_db -a -v ON_ERROR_STOP=1 -f ./tests/integration_tests/mssql_source/assertions/large_table_test_full_refresh_assertions.sql
 if %errorlevel% neq 0 exit /b %errorlevel%
+
+psql -U postgres -d rdl_integration_test_target_db -a -v ON_ERROR_STOP=1 -f ./tests/integration_tests/mssql_source/assertions/compound_pk_test_full_refresh_assertions.sql
+if %errorlevel% neq 0 exit /b %errorlevel%

tests/integration_tests/test_incremental_refresh_from_mssql.cmd

@@ -4,3 +4,6 @@ if %errorlevel% neq 0 exit /b %errorlevel%
 
 psql -U postgres -d rdl_integration_test_target_db -a -v ON_ERROR_STOP=1 -f ./tests/integration_tests/mssql_source/assertions/large_table_test_full_refresh_assertions.sql
 if %errorlevel% neq 0 exit /b %errorlevel%
+
+psql -U postgres -d rdl_integration_test_target_db -a -v ON_ERROR_STOP=1 -f ./tests/integration_tests/mssql_source/assertions/compound_pk_test_full_refresh_assertions.sql
+if %errorlevel% neq 0 exit /b %errorlevel%