sql-testing-library 0.11.0__tar.gz → 0.13.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (23) hide show
  1. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/CHANGELOG.md +17 -0
  2. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/PKG-INFO +136 -13
  3. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/README.md +135 -12
  4. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/pyproject.toml +4 -1
  5. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/presto.py +20 -6
  6. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/snowflake.py +99 -9
  7. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_core.py +23 -2
  8. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_pytest_plugin.py +33 -10
  9. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_sql_logger.py +45 -7
  10. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_sql_utils.py +152 -2
  11. sql_testing_library-0.13.0/src/sql_testing_library/_types.py +412 -0
  12. sql_testing_library-0.11.0/src/sql_testing_library/_types.py +0 -203
  13. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/LICENSE +0 -0
  14. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/__init__.py +0 -0
  15. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/__init__.py +0 -0
  16. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/athena.py +0 -0
  17. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/base.py +0 -0
  18. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/bigquery.py +0 -0
  19. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/redshift.py +0 -0
  20. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_adapters/trino.py +0 -0
  21. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_exceptions.py +0 -0
  22. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/_mock_table.py +0 -0
  23. {sql_testing_library-0.11.0 → sql_testing_library-0.13.0}/src/sql_testing_library/py.typed +0 -0
@@ -5,6 +5,23 @@ All notable changes to this project will be documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## 0.13.0 (2025-06-27)
9
+
10
+ ### Feat
11
+
12
+ - add pytest-xdist support for parallel test execution (#105)
13
+
14
+ ### Fix
15
+
16
+ - **snowflake**: fix issue related to physical view for snowflake (#104)
17
+
18
+ ## 0.12.0 (2025-06-25)
19
+
20
+ ### Feat
21
+
22
+ - **athena/trino**: add support for struct/ROW types (#102)
23
+ - **snowflake**: add key-pair authentication for MFA support (#103)
24
+
8
25
  ## 0.11.0 (2025-06-16)
9
26
 
10
27
  ### Feat
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: sql-testing-library
3
- Version: 0.11.0
3
+ Version: 0.13.0
4
4
  Summary: A powerful Python framework for unit testing SQL queries across BigQuery, Snowflake, Redshift, Athena, and Trino with mock data
5
5
  License: MIT
6
6
  Keywords: sql,testing,unit-testing,mock-data,database-testing,bigquery,snowflake,redshift,athena,trino,data-engineering,etl-testing,sql-validation,query-testing
@@ -137,16 +137,16 @@ The library supports different data types across database engines. All checkmark
137
137
  | **Decimal Array** | `List[Decimal]` | ✅ | ✅ | ✅ | ✅ | ✅ |
138
138
  | **Optional Array** | `Optional[List[T]]` | ✅ | ✅ | ✅ | ✅ | ✅ |
139
139
  | **Map/Dict** | `Dict[K, V]` | ✅ | ✅ | ✅ | ✅ | ✅ |
140
- | **Struct/Record** | `dataclass` | ❌ | | ❌ | | ❌ |
140
+ | **Struct/Record** | `dataclass` | ❌ | | ❌ | | ❌ |
141
141
  | **Nested Arrays** | `List[List[T]]` | ❌ | ❌ | ❌ | ❌ | ❌ |
142
142
 
143
143
  ### Database-Specific Notes
144
144
 
145
- - **BigQuery**: NULL arrays become empty arrays `[]`; uses scientific notation for large decimals; dict/map types stored as JSON strings
146
- - **Athena**: 256KB query size limit; supports arrays and maps using `ARRAY[]` and `MAP(ARRAY[], ARRAY[])` syntax
147
- - **Redshift**: Arrays and maps implemented via SUPER type (JSON parsing); 16MB query size limit
148
- - **Trino**: Memory catalog for testing; excellent decimal precision; supports arrays and maps
149
- - **Snowflake**: Column names normalized to lowercase; 1MB query size limit; dict/map types implemented via VARIANT type (JSON parsing)
145
+ - **BigQuery**: NULL arrays become empty arrays `[]`; uses scientific notation for large decimals; dict/map types stored as JSON strings; struct types not yet supported
146
+ - **Athena**: 256KB query size limit; supports arrays and maps using `ARRAY[]` and `MAP(ARRAY[], ARRAY[])` syntax; supports struct types using `ROW` with named fields (dataclasses and Pydantic models)
147
+ - **Redshift**: Arrays and maps implemented via SUPER type (JSON parsing); 16MB query size limit; struct types not yet supported
148
+ - **Trino**: Memory catalog for testing; excellent decimal precision; supports arrays, maps, and struct types using `ROW` with named fields (dataclasses and Pydantic models)
149
+ - **Snowflake**: Column names normalized to lowercase; 1MB query size limit; dict/map types implemented via VARIANT type (JSON parsing); struct types not yet supported
150
150
 
151
151
  ## Execution Modes Support
152
152
 
@@ -155,7 +155,7 @@ The library supports two execution modes for mock data injection. **CTE Mode is
155
155
  | Execution Mode | Description | BigQuery | Athena | Redshift | Trino | Snowflake |
156
156
  |----------------|-------------|----------|--------|----------|-------|-----------|
157
157
  | **CTE Mode** | Mock data injected as Common Table Expressions | ✅ | ✅ | ✅ | ✅ | ✅ |
158
- | **Physical Tables** | Mock data created as temporary tables | ✅ | ✅ | ✅ | ✅ | ⚠️* |
158
+ | **Physical Tables** | Mock data created as temporary tables | ✅ | ✅ | ✅ | ✅ | |
159
159
 
160
160
  ### Execution Mode Details
161
161
 
@@ -306,7 +306,7 @@ def test_physical_tables():
306
306
  **Notes:**
307
307
  - **CTE Mode**: Default mode, works with all database engines, suitable for most use cases
308
308
  - **Physical Tables**: Used automatically when CTE queries exceed database size limits or when explicitly requested
309
- - **⚠️ Snowflake Physical Tables**: Currently disabled in test environment due to temporary table visibility limitations. CTE mode works perfectly for all use cases.
309
+ - **Snowflake**: Full support for both CTE and physical table modes
310
310
 
311
311
  ## Installation
312
312
 
@@ -399,11 +399,18 @@ credentials_path = <path to credentials json>
399
399
  # [sql_testing.snowflake]
400
400
  # account = <account-identifier>
401
401
  # user = <snowflake_user>
402
- # password = <snowflake_password>
403
402
  # database = <test_database>
404
403
  # schema = <PUBLIC> # Optional: default schema is 'PUBLIC'
405
404
  # warehouse = <compute_wh> # Required: specify a warehouse
406
405
  # role = <role_name> # Optional: specify a role
406
+ #
407
+ # # Authentication (choose one):
408
+ # # Option 1: Key-pair authentication (recommended for MFA)
409
+ # private_key_path = </path/to/private_key.pem>
410
+ # # Or use environment variable SNOWFLAKE_PRIVATE_KEY
411
+ #
412
+ # # Option 2: Password authentication (for accounts without MFA)
413
+ # password = <snowflake_password>
407
414
  ```
408
415
 
409
416
  ### Database Context Understanding
@@ -565,6 +572,125 @@ def test_pattern_3():
565
572
  )
566
573
  ```
567
574
 
575
+ ### Working with Struct Types (Athena and Trino)
576
+
577
+ The library supports struct/record types using Python dataclasses or Pydantic models for Athena and Trino:
578
+
579
+ ```python
580
+ from dataclasses import dataclass
581
+ from decimal import Decimal
582
+ from pydantic import BaseModel
583
+ from sql_testing_library import sql_test, TestCase
584
+ from sql_testing_library.mock_table import BaseMockTable
585
+
586
+ # Define nested structs using dataclasses
587
+ @dataclass
588
+ class Address:
589
+ street: str
590
+ city: str
591
+ state: str
592
+ zip_code: str
593
+
594
+ @dataclass
595
+ class Employee:
596
+ id: int
597
+ name: str
598
+ salary: Decimal
599
+ address: Address # Nested struct
600
+ is_active: bool = True
601
+
602
+ # Or use Pydantic models
603
+ class AddressPydantic(BaseModel):
604
+ street: str
605
+ city: str
606
+ state: str
607
+ zip_code: str
608
+
609
+ class EmployeeResultPydantic(BaseModel):
610
+ id: int
611
+ name: str
612
+ city: str # Extracted from nested struct
613
+
614
+ # Mock table with struct data
615
+ class EmployeesMockTable(BaseMockTable):
616
+ def get_database_name(self) -> str:
617
+ return "test_db"
618
+
619
+ def get_table_name(self) -> str:
620
+ return "employees"
621
+
622
+ # Test with struct types
623
+ @sql_test(
624
+ adapter_type="athena", # or "trino"
625
+ mock_tables=[
626
+ EmployeesMockTable([
627
+ Employee(
628
+ id=1,
629
+ name="Alice Johnson",
630
+ salary=Decimal("120000.00"),
631
+ address=Address(
632
+ street="123 Tech Lane",
633
+ city="San Francisco",
634
+ state="CA",
635
+ zip_code="94105"
636
+ ),
637
+ is_active=True
638
+ ),
639
+ Employee(
640
+ id=2,
641
+ name="Bob Smith",
642
+ salary=Decimal("95000.00"),
643
+ address=Address(
644
+ street="456 Oak Ave",
645
+ city="New York",
646
+ state="NY",
647
+ zip_code="10001"
648
+ ),
649
+ is_active=False
650
+ )
651
+ ])
652
+ ],
653
+ result_class=EmployeeResultPydantic
654
+ )
655
+ def test_struct_with_dot_notation():
656
+ return TestCase(
657
+ query="""
658
+ SELECT
659
+ id,
660
+ name,
661
+ address.city as city -- Access nested field with dot notation
662
+ FROM employees
663
+ WHERE address.state = 'CA' -- Use struct fields in WHERE clause
664
+ """,
665
+ default_namespace="test_db"
666
+ )
667
+
668
+ # You can also query entire structs
669
+ @sql_test(
670
+ adapter_type="trino",
671
+ mock_tables=[EmployeesMockTable([...])],
672
+ result_class=dict # Returns full struct as dict
673
+ )
674
+ def test_query_full_struct():
675
+ return TestCase(
676
+ query="SELECT id, name, address FROM employees",
677
+ default_namespace="test_db"
678
+ )
679
+ ```
680
+
681
+ **Struct Type Features:**
682
+ - **Nested Structures**: Support for deeply nested structs using dataclasses or Pydantic models
683
+ - **Dot Notation**: Access struct fields using `struct.field` syntax in queries
684
+ - **Type Safety**: Full type conversion between Python objects and SQL ROW types
685
+ - **NULL Handling**: Proper handling of optional struct fields
686
+ - **WHERE Clause**: Use struct fields in filtering conditions
687
+ - **List of Structs**: Full support for `List[StructType]` with array operations
688
+
689
+ **SQL Type Mapping:**
690
+ - Python dataclass/Pydantic model → SQL `ROW(field1 type1, field2 type2, ...)`
691
+ - Nested structs are fully supported
692
+ - All struct values are properly cast to ensure type consistency
693
+
568
694
  3. **Run with pytest**:
569
695
 
570
696
  ```bash
@@ -1046,9 +1172,6 @@ For detailed usage and configuration options, see the example files included.
1046
1172
 
1047
1173
  The library has a few known limitations that are planned to be addressed in future updates:
1048
1174
 
1049
- ### Snowflake Support
1050
- - Physical table tests for Snowflake are currently skipped due to complex mocking requirements
1051
-
1052
1175
  ### General Improvements
1053
1176
  - Add support for more SQL dialects
1054
1177
  - Improve error handling for malformed SQL
@@ -80,16 +80,16 @@ The library supports different data types across database engines. All checkmark
80
80
  | **Decimal Array** | `List[Decimal]` | ✅ | ✅ | ✅ | ✅ | ✅ |
81
81
  | **Optional Array** | `Optional[List[T]]` | ✅ | ✅ | ✅ | ✅ | ✅ |
82
82
  | **Map/Dict** | `Dict[K, V]` | ✅ | ✅ | ✅ | ✅ | ✅ |
83
- | **Struct/Record** | `dataclass` | ❌ | | ❌ | | ❌ |
83
+ | **Struct/Record** | `dataclass` | ❌ | | ❌ | | ❌ |
84
84
  | **Nested Arrays** | `List[List[T]]` | ❌ | ❌ | ❌ | ❌ | ❌ |
85
85
 
86
86
  ### Database-Specific Notes
87
87
 
88
- - **BigQuery**: NULL arrays become empty arrays `[]`; uses scientific notation for large decimals; dict/map types stored as JSON strings
89
- - **Athena**: 256KB query size limit; supports arrays and maps using `ARRAY[]` and `MAP(ARRAY[], ARRAY[])` syntax
90
- - **Redshift**: Arrays and maps implemented via SUPER type (JSON parsing); 16MB query size limit
91
- - **Trino**: Memory catalog for testing; excellent decimal precision; supports arrays and maps
92
- - **Snowflake**: Column names normalized to lowercase; 1MB query size limit; dict/map types implemented via VARIANT type (JSON parsing)
88
+ - **BigQuery**: NULL arrays become empty arrays `[]`; uses scientific notation for large decimals; dict/map types stored as JSON strings; struct types not yet supported
89
+ - **Athena**: 256KB query size limit; supports arrays and maps using `ARRAY[]` and `MAP(ARRAY[], ARRAY[])` syntax; supports struct types using `ROW` with named fields (dataclasses and Pydantic models)
90
+ - **Redshift**: Arrays and maps implemented via SUPER type (JSON parsing); 16MB query size limit; struct types not yet supported
91
+ - **Trino**: Memory catalog for testing; excellent decimal precision; supports arrays, maps, and struct types using `ROW` with named fields (dataclasses and Pydantic models)
92
+ - **Snowflake**: Column names normalized to lowercase; 1MB query size limit; dict/map types implemented via VARIANT type (JSON parsing); struct types not yet supported
93
93
 
94
94
  ## Execution Modes Support
95
95
 
@@ -98,7 +98,7 @@ The library supports two execution modes for mock data injection. **CTE Mode is
98
98
  | Execution Mode | Description | BigQuery | Athena | Redshift | Trino | Snowflake |
99
99
  |----------------|-------------|----------|--------|----------|-------|-----------|
100
100
  | **CTE Mode** | Mock data injected as Common Table Expressions | ✅ | ✅ | ✅ | ✅ | ✅ |
101
- | **Physical Tables** | Mock data created as temporary tables | ✅ | ✅ | ✅ | ✅ | ⚠️* |
101
+ | **Physical Tables** | Mock data created as temporary tables | ✅ | ✅ | ✅ | ✅ | |
102
102
 
103
103
  ### Execution Mode Details
104
104
 
@@ -249,7 +249,7 @@ def test_physical_tables():
249
249
  **Notes:**
250
250
  - **CTE Mode**: Default mode, works with all database engines, suitable for most use cases
251
251
  - **Physical Tables**: Used automatically when CTE queries exceed database size limits or when explicitly requested
252
- - **⚠️ Snowflake Physical Tables**: Currently disabled in test environment due to temporary table visibility limitations. CTE mode works perfectly for all use cases.
252
+ - **Snowflake**: Full support for both CTE and physical table modes
253
253
 
254
254
  ## Installation
255
255
 
@@ -342,11 +342,18 @@ credentials_path = <path to credentials json>
342
342
  # [sql_testing.snowflake]
343
343
  # account = <account-identifier>
344
344
  # user = <snowflake_user>
345
- # password = <snowflake_password>
346
345
  # database = <test_database>
347
346
  # schema = <PUBLIC> # Optional: default schema is 'PUBLIC'
348
347
  # warehouse = <compute_wh> # Required: specify a warehouse
349
348
  # role = <role_name> # Optional: specify a role
349
+ #
350
+ # # Authentication (choose one):
351
+ # # Option 1: Key-pair authentication (recommended for MFA)
352
+ # private_key_path = </path/to/private_key.pem>
353
+ # # Or use environment variable SNOWFLAKE_PRIVATE_KEY
354
+ #
355
+ # # Option 2: Password authentication (for accounts without MFA)
356
+ # password = <snowflake_password>
350
357
  ```
351
358
 
352
359
  ### Database Context Understanding
@@ -508,6 +515,125 @@ def test_pattern_3():
508
515
  )
509
516
  ```
510
517
 
518
+ ### Working with Struct Types (Athena and Trino)
519
+
520
+ The library supports struct/record types using Python dataclasses or Pydantic models for Athena and Trino:
521
+
522
+ ```python
523
+ from dataclasses import dataclass
524
+ from decimal import Decimal
525
+ from pydantic import BaseModel
526
+ from sql_testing_library import sql_test, TestCase
527
+ from sql_testing_library.mock_table import BaseMockTable
528
+
529
+ # Define nested structs using dataclasses
530
+ @dataclass
531
+ class Address:
532
+ street: str
533
+ city: str
534
+ state: str
535
+ zip_code: str
536
+
537
+ @dataclass
538
+ class Employee:
539
+ id: int
540
+ name: str
541
+ salary: Decimal
542
+ address: Address # Nested struct
543
+ is_active: bool = True
544
+
545
+ # Or use Pydantic models
546
+ class AddressPydantic(BaseModel):
547
+ street: str
548
+ city: str
549
+ state: str
550
+ zip_code: str
551
+
552
+ class EmployeeResultPydantic(BaseModel):
553
+ id: int
554
+ name: str
555
+ city: str # Extracted from nested struct
556
+
557
+ # Mock table with struct data
558
+ class EmployeesMockTable(BaseMockTable):
559
+ def get_database_name(self) -> str:
560
+ return "test_db"
561
+
562
+ def get_table_name(self) -> str:
563
+ return "employees"
564
+
565
+ # Test with struct types
566
+ @sql_test(
567
+ adapter_type="athena", # or "trino"
568
+ mock_tables=[
569
+ EmployeesMockTable([
570
+ Employee(
571
+ id=1,
572
+ name="Alice Johnson",
573
+ salary=Decimal("120000.00"),
574
+ address=Address(
575
+ street="123 Tech Lane",
576
+ city="San Francisco",
577
+ state="CA",
578
+ zip_code="94105"
579
+ ),
580
+ is_active=True
581
+ ),
582
+ Employee(
583
+ id=2,
584
+ name="Bob Smith",
585
+ salary=Decimal("95000.00"),
586
+ address=Address(
587
+ street="456 Oak Ave",
588
+ city="New York",
589
+ state="NY",
590
+ zip_code="10001"
591
+ ),
592
+ is_active=False
593
+ )
594
+ ])
595
+ ],
596
+ result_class=EmployeeResultPydantic
597
+ )
598
+ def test_struct_with_dot_notation():
599
+ return TestCase(
600
+ query="""
601
+ SELECT
602
+ id,
603
+ name,
604
+ address.city as city -- Access nested field with dot notation
605
+ FROM employees
606
+ WHERE address.state = 'CA' -- Use struct fields in WHERE clause
607
+ """,
608
+ default_namespace="test_db"
609
+ )
610
+
611
+ # You can also query entire structs
612
+ @sql_test(
613
+ adapter_type="trino",
614
+ mock_tables=[EmployeesMockTable([...])],
615
+ result_class=dict # Returns full struct as dict
616
+ )
617
+ def test_query_full_struct():
618
+ return TestCase(
619
+ query="SELECT id, name, address FROM employees",
620
+ default_namespace="test_db"
621
+ )
622
+ ```
623
+
624
+ **Struct Type Features:**
625
+ - **Nested Structures**: Support for deeply nested structs using dataclasses or Pydantic models
626
+ - **Dot Notation**: Access struct fields using `struct.field` syntax in queries
627
+ - **Type Safety**: Full type conversion between Python objects and SQL ROW types
628
+ - **NULL Handling**: Proper handling of optional struct fields
629
+ - **WHERE Clause**: Use struct fields in filtering conditions
630
+ - **List of Structs**: Full support for `List[StructType]` with array operations
631
+
632
+ **SQL Type Mapping:**
633
+ - Python dataclass/Pydantic model → SQL `ROW(field1 type1, field2 type2, ...)`
634
+ - Nested structs are fully supported
635
+ - All struct values are properly cast to ensure type consistency
636
+
511
637
  3. **Run with pytest**:
512
638
 
513
639
  ```bash
@@ -989,9 +1115,6 @@ For detailed usage and configuration options, see the example files included.
989
1115
 
990
1116
  The library has a few known limitations that are planned to be addressed in future updates:
991
1117
 
992
- ### Snowflake Support
993
- - Physical table tests for Snowflake are currently skipped due to complex mocking requirements
994
-
995
1118
  ### General Improvements
996
1119
  - Add support for more SQL dialects
997
1120
  - Improve error handling for malformed SQL
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
4
4
 
5
5
  [tool.poetry]
6
6
  name = "sql-testing-library"
7
- version = "0.11.0"
7
+ version = "0.13.0"
8
8
  description = "A powerful Python framework for unit testing SQL queries across BigQuery, Snowflake, Redshift, Athena, and Trino with mock data"
9
9
  authors = ["Gurmeet Saran <gurmeetx@gmail.com>", "Kushal Thakkar <kushal.thakkar@gmail.com>"]
10
10
  maintainers = ["Gurmeet Saran <gurmeetx@gmail.com>", "Kushal Thakkar <kushal.thakkar@gmail.com>"]
@@ -104,6 +104,7 @@ pytest = ">=7.0.0"
104
104
  pytest-asyncio = ">=0.21.0"
105
105
  pytest-mock = ">=3.10.0"
106
106
  pytest-cov = ">=4.1.0"
107
+ pytest-xdist = ">=3.0.0"
107
108
  pytest-rerunfailures = ">=12.0.0"
108
109
  black = ">=22.0.0"
109
110
  flake8 = ">=4.0.0"
@@ -193,6 +194,8 @@ unfixable = ["T201", "T203"]
193
194
  # Allow print statements in test files for debugging and output
194
195
  "tests/test_*.py" = ["T201"]
195
196
  "tests/*/test_*.py" = ["T201"]
197
+ # Allow print statements in example files
198
+ "examples/*.py" = ["T201"]
196
199
 
197
200
  [tool.ruff.lint.isort]
198
201
  lines-after-imports = 2
@@ -6,7 +6,7 @@ from decimal import Decimal
6
6
  from typing import Any, List, Tuple, Type, Union, get_args
7
7
 
8
8
  from .._mock_table import BaseMockTable
9
- from .._types import BaseTypeConverter
9
+ from .._types import BaseTypeConverter, is_struct_type
10
10
  from .base import DatabaseAdapter
11
11
 
12
12
 
@@ -78,6 +78,8 @@ class PrestoBaseAdapter(DatabaseAdapter):
78
78
 
79
79
  def _get_sql_type(self, python_type: Type) -> str:
80
80
  """Convert Python type to SQL type string."""
81
+ from .._sql_utils import get_sql_type_string
82
+
81
83
  # Handle Optional types
82
84
  if hasattr(python_type, "__origin__") and python_type.__origin__ is Union:
83
85
  # Extract the non-None type from Optional[T]
@@ -85,10 +87,14 @@ class PrestoBaseAdapter(DatabaseAdapter):
85
87
  if non_none_types:
86
88
  python_type = non_none_types[0]
87
89
 
90
+ # Handle struct types (dataclass or Pydantic model)
91
+ if is_struct_type(python_type):
92
+ return get_sql_type_string(python_type, self.get_sqlglot_dialect())
93
+
88
94
  # Handle List types
89
95
  if hasattr(python_type, "__origin__") and python_type.__origin__ is list:
90
96
  element_type = get_args(python_type)[0] if get_args(python_type) else str
91
- element_sql_type = self.TYPE_MAPPING.get(element_type, "VARCHAR")
97
+ element_sql_type = self._get_sql_type(element_type) # Recursive call for nested types
92
98
  return f"ARRAY({element_sql_type})"
93
99
 
94
100
  # Handle Dict/Map types
@@ -96,12 +102,20 @@ class PrestoBaseAdapter(DatabaseAdapter):
96
102
  type_args = get_args(python_type)
97
103
  key_type = type_args[0] if type_args else str
98
104
  value_type = type_args[1] if len(type_args) > 1 else str
99
- key_sql_type = self.TYPE_MAPPING.get(key_type, "VARCHAR")
100
- value_sql_type = self.TYPE_MAPPING.get(value_type, "VARCHAR")
105
+ key_sql_type = self._get_sql_type(key_type) # Recursive call for nested types
106
+ value_sql_type = self._get_sql_type(value_type) # Recursive call for nested types
101
107
  return f"MAP({key_sql_type}, {value_sql_type})"
102
108
 
103
- # Regular types
104
- return self.TYPE_MAPPING.get(python_type, "VARCHAR")
109
+ # Regular types - use the mapping
110
+ dialect = self.get_sqlglot_dialect()
111
+ return get_sql_type_string(python_type, dialect)
112
+
113
+ def _get_struct_sql_type(self, struct_type: Type) -> str:
114
+ """Get SQL ROW type definition for a struct (dataclass or Pydantic model)."""
115
+ from .._sql_utils import get_sql_type_string
116
+
117
+ dialect = self.get_sqlglot_dialect()
118
+ return get_sql_type_string(struct_type, dialect)
105
119
 
106
120
  def _generate_column_definitions(self, column_types: dict) -> str:
107
121
  """Generate SQL column definitions from column types."""
@@ -1,10 +1,11 @@
1
1
  """Snowflake adapter implementation."""
2
2
 
3
3
  import logging
4
+ import os
4
5
  import time
5
6
  from datetime import date, datetime
6
7
  from decimal import Decimal
7
- from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union, get_args
8
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, get_args
8
9
 
9
10
 
10
11
  if TYPE_CHECKING:
@@ -67,11 +68,13 @@ class SnowflakeAdapter(DatabaseAdapter):
67
68
  self,
68
69
  account: str,
69
70
  user: str,
70
- password: str,
71
- database: str,
71
+ password: Optional[str] = None,
72
+ database: str = "",
72
73
  schema: str = "PUBLIC",
73
74
  warehouse: Optional[str] = None,
74
75
  role: Optional[str] = None,
76
+ private_key_path: Optional[str] = None,
77
+ private_key_passphrase: Optional[str] = None,
75
78
  ) -> None:
76
79
  if not has_snowflake:
77
80
  raise ImportError(
@@ -88,7 +91,77 @@ class SnowflakeAdapter(DatabaseAdapter):
88
91
  self.schema = schema
89
92
  self.warehouse = warehouse
90
93
  self.role = role
94
+ self.private_key_path = private_key_path
95
+ self.private_key_passphrase = private_key_passphrase
91
96
  self.conn: Optional[Any] = None
97
+ self._private_key: Optional[bytes] = None
98
+
99
+ def _load_private_key(self) -> bytes:
100
+ """Load private key from file or environment variable and convert to DER format."""
101
+ if self._private_key:
102
+ return self._private_key
103
+
104
+ # Try to load from file path
105
+ if self.private_key_path and os.path.exists(self.private_key_path):
106
+ with open(self.private_key_path, "rb") as key_file:
107
+ private_key_data = key_file.read()
108
+ # Try to load from environment variable
109
+ elif os.environ.get("SNOWFLAKE_PRIVATE_KEY"):
110
+ private_key_data = os.environ["SNOWFLAKE_PRIVATE_KEY"].encode()
111
+ else:
112
+ raise ValueError(
113
+ "Private key not found. Provide private_key_path or "
114
+ "set SNOWFLAKE_PRIVATE_KEY environment variable"
115
+ )
116
+
117
+ # Import cryptography modules
118
+ from cryptography.hazmat.backends import default_backend
119
+ from cryptography.hazmat.primitives import serialization
120
+ from cryptography.hazmat.primitives.serialization import (
121
+ load_der_private_key,
122
+ load_pem_private_key,
123
+ )
124
+
125
+ # Determine if we have a passphrase
126
+ passphrase_str = self.private_key_passphrase or os.environ.get(
127
+ "SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"
128
+ )
129
+ passphrase = passphrase_str.encode() if passphrase_str else None
130
+
131
+ try:
132
+ # Check if it's already in DER format
133
+ if not private_key_data.startswith(b"-----"):
134
+ # Already in DER format, try to load it
135
+ try:
136
+ private_key_obj = load_der_private_key(
137
+ private_key_data, password=passphrase, backend=default_backend()
138
+ )
139
+ private_key = private_key_data
140
+ except Exception:
141
+ # Not valid DER, treat as PEM
142
+ pass
143
+
144
+ # Load PEM private key (handles both PKCS#1 and PKCS#8)
145
+ if private_key_data.startswith(b"-----"):
146
+ private_key_obj = load_pem_private_key(
147
+ private_key_data, password=passphrase, backend=default_backend()
148
+ )
149
+
150
+ # Convert to DER format (PKCS#8) for Snowflake
151
+ private_key = private_key_obj.private_bytes(
152
+ encoding=serialization.Encoding.DER,
153
+ format=serialization.PrivateFormat.PKCS8,
154
+ encryption_algorithm=serialization.NoEncryption(),
155
+ )
156
+ except Exception as e:
157
+ raise ValueError(
158
+ f"Failed to load private key: {e}. "
159
+ "Ensure the key is in PEM format (PKCS#1 or PKCS#8) "
160
+ "and the passphrase (if any) is correct."
161
+ ) from e
162
+
163
+ self._private_key = private_key
164
+ return private_key
92
165
 
93
166
  def _get_connection(self) -> Any:
94
167
  """Get or create a connection to Snowflake."""
@@ -96,14 +169,31 @@ class SnowflakeAdapter(DatabaseAdapter):
96
169
 
97
170
  # Create a new connection if needed
98
171
  if self.conn is None:
99
- conn_params = {
172
+ conn_params: Dict[str, Any] = {
100
173
  "account": self.account,
101
174
  "user": self.user,
102
- "password": self.password,
103
- "database": self.database,
104
- "schema": self.schema,
105
175
  }
106
176
 
177
+ # Add optional parameters
178
+ if self.database:
179
+ conn_params["database"] = self.database
180
+
181
+ if self.schema:
182
+ conn_params["schema"] = self.schema
183
+
184
+ # Handle authentication
185
+ if self.private_key_path or os.environ.get("SNOWFLAKE_PRIVATE_KEY"):
186
+ # Use key-pair authentication
187
+ conn_params["private_key"] = self._load_private_key()
188
+ elif self.password:
189
+ conn_params["password"] = self.password
190
+ else:
191
+ raise ValueError(
192
+ "No authentication method provided. Please provide one of: "
193
+ "1) password, 2) private_key_path, or "
194
+ "3) set SNOWFLAKE_PRIVATE_KEY environment variable"
195
+ )
196
+
107
197
  if self.warehouse:
108
198
  conn_params["warehouse"] = self.warehouse
109
199
 
@@ -279,7 +369,7 @@ class SnowflakeAdapter(DatabaseAdapter):
279
369
  snowflake_type = "VARIANT"
280
370
  else:
281
371
  snowflake_type = type_mapping.get(col_type, "VARCHAR")
282
- column_defs.append(f'"{col_name}" {snowflake_type}')
372
+ column_defs.append(f"{col_name} {snowflake_type}")
283
373
 
284
374
  columns_sql = ",\n ".join(column_defs)
285
375
 
@@ -300,7 +390,7 @@ class SnowflakeAdapter(DatabaseAdapter):
300
390
  col_type = column_types.get(col_name, str)
301
391
  value = first_row[col_name]
302
392
  formatted_value = self.format_value_for_cte(value, col_type)
303
- select_expressions.append(f'{formatted_value} AS "{col_name}"')
393
+ select_expressions.append(f"{formatted_value} AS {col_name}")
304
394
 
305
395
  # Start with the first row in the SELECT
306
396
  select_sql = f"SELECT {', '.join(select_expressions)}"