Skip to content

Commit 45dfe66

Browse files
Brooke-whiteBrooke-white
committed
feat(datatype): support geography
1 parent 03b9d6d commit 45dfe66

File tree

8 files changed

+56
-4
lines changed

8 files changed

+56
-4
lines changed

‎README.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,8 @@ Incoming data from Amazon Redshift is treated as follows:
400400
+--------------------------+-------------------+
401401
| VARBYTE | bytes |
402402
+--------------------------+-------------------+
403+
| GEOGRAPHY | str |
404+
+--------------------------+-------------------+
403405

404406
Logging
405407
~~~~~~~~~~~~

‎redshift_connector/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
DECIMAL_ARRAY,
5959
FLOAT,
6060
FLOAT_ARRAY,
61+
GEOGRAPHY,
6162
GEOMETRY,
6263
GEOMETRYHEX,
6364
INET,
@@ -421,6 +422,7 @@ def connect(
421422
"DECIMAL_ARRAY",
422423
"FLOAT",
423424
"FLOAT_ARRAY",
425+
"GEOGRAPHY",
424426
"GEOMETRY",
425427
"GEOMETRYHEX",
426428
"INET",
@@ -448,6 +450,7 @@ def connect(
448450
"TIMETZ",
449451
"UNKNOWN",
450452
"UUID_TYPE",
453+
"VARBYTE",
451454
"VARCHAR",
452455
"VARCHAR_ARRAY",
453456
"XID",

‎redshift_connector/core.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
date_in,
6464
date_recv_binary,
6565
float_array_recv,
66+
geographyhex_recv,
6667
h_pack,
6768
h_unpack,
6869
i_pack,
@@ -90,6 +91,7 @@
9091
from redshift_connector.utils.type_utils import (
9192
BIGINT,
9293
DATE,
94+
GEOGRAPHY,
9395
INTEGER,
9496
INTEGER_ARRAY,
9597
NUMERIC,
@@ -685,6 +687,7 @@ def _enable_protocol_based_conversion_funcs(self: "Connection"):
685687
if self._client_protocol_version >= ClientProtocolVersion.BINARY.value:
686688
self.pg_types[NUMERIC] = (FC_BINARY, numeric_in_binary)
687689
self.pg_types[DATE] = (FC_BINARY, date_recv_binary)
690+
self.pg_types[GEOGRAPHY] = (FC_BINARY, geographyhex_recv) # GEOGRAPHY
688691
self.pg_types[TIME] = (FC_BINARY, time_recv_binary)
689692
self.pg_types[TIMETZ] = (FC_BINARY, timetz_recv_binary)
690693
self.pg_types[1002] = (FC_BINARY, array_recv_binary) # CHAR[]
@@ -700,6 +703,7 @@ def _enable_protocol_based_conversion_funcs(self: "Connection"):
700703
self.pg_types[NUMERIC] = (FC_TEXT, numeric_in)
701704
self.pg_types[TIME] = (FC_TEXT, time_in)
702705
self.pg_types[DATE] = (FC_TEXT, date_in)
706+
self.pg_types[GEOGRAPHY] = (FC_TEXT, text_recv) # GEOGRAPHY
703707
self.pg_types[TIMETZ] = (FC_BINARY, timetz_recv_binary)
704708
self.pg_types[1002] = (FC_TEXT, array_recv_text) # CHAR[]
705709
self.pg_types[SMALLINT_ARRAY] = (FC_TEXT, int_array_recv) # INT2[]

‎redshift_connector/cursor.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1174,6 +1174,7 @@ def __build_local_schema_columns_query(
11741174
"when 'geometry' THEN -4 "
11751175
"when 'super' THEN -16 "
11761176
"when 'varbyte' THEN -4 "
1177+
"when 'geography' THEN -4 "
11771178
"else 1111 END as SMALLINT) AS DATA_TYPE, "
11781179
"t.typname as TYPE_NAME, "
11791180
"case typname "
@@ -1211,6 +1212,7 @@ def __build_local_schema_columns_query(
12111212
"when 'geometry' THEN NULL "
12121213
"when 'super' THEN NULL "
12131214
"when 'varbyte' THEN NULL "
1215+
"when 'geography' THEN NULL "
12141216
"else 2147483647 end as COLUMN_SIZE , "
12151217
"null as BUFFER_LENGTH , "
12161218
"case typname "
@@ -1221,6 +1223,7 @@ def __build_local_schema_columns_query(
12211223
"when 'geometry' then NULL "
12221224
"when 'super' then NULL "
12231225
"when 'varbyte' then NULL "
1226+
"when 'geography' then NULL "
12241227
"else 0 end as DECIMAL_DIGITS, "
12251228
"10 AS NUM_PREC_RADIX , "
12261229
"case a.attnotnull OR (t.typtype = 'd' AND t.typnotnull) "
@@ -1273,6 +1276,7 @@ def __build_local_schema_columns_query(
12731276
"when 'geometry' THEN -4 "
12741277
"when 'super' THEN -16 "
12751278
"when 'varbyte' THEN -4 "
1279+
"when 'geography' THEN -4 "
12761280
"else 1111 END as SMALLINT) AS SQL_DATA_TYPE, "
12771281
"CAST(NULL AS SMALLINT) as SQL_DATETIME_SUB , "
12781282
"case typname "
@@ -1310,6 +1314,7 @@ def __build_local_schema_columns_query(
13101314
"when 'geometry' THEN NULL "
13111315
"when 'super' THEN NULL "
13121316
"when 'varbyte' THEN NULL "
1317+
"when 'geography' THEN NULL "
13131318
"else 2147483647 end as CHAR_OCTET_LENGTH , "
13141319
"a.attnum AS ORDINAL_POSITION, "
13151320
"case a.attnotnull OR (t.typtype = 'd' AND t.typnotnull) "
@@ -1391,6 +1396,7 @@ def __build_local_schema_columns_query(
13911396
"WHEN 'geometry' THEN -4 "
13921397
"WHEN 'super' THEN -16 "
13931398
"WHEN 'varbyte' THEN -4 "
1399+
"WHEN 'geography' THEN -4"
13941400
"ELSE 1111 END AS SMALLINT) AS DATA_TYPE, "
13951401
"COALESCE(NULL,CASE columntype WHEN 'boolean' THEN 'bool' "
13961402
"WHEN 'character varying' THEN 'varchar' "
@@ -1441,6 +1447,7 @@ def __build_local_schema_columns_query(
14411447
"WHEN 'geometry' THEN NULL "
14421448
"WHEN 'super' THEN NULL "
14431449
"WHEN 'varbyte' THEN NULL "
1450+
"WHEN 'geography' THEN NULL "
14441451
"ELSE 2147483647 END AS COLUMN_SIZE, "
14451452
"NULL AS BUFFER_LENGTH, "
14461453
"CASE columntype "
@@ -1453,6 +1460,7 @@ def __build_local_schema_columns_query(
14531460
"WHEN 'geometry' THEN NULL "
14541461
"WHEN 'super' THEN NULL "
14551462
"WHEN 'varbyte' THEN NULL "
1463+
"WHEN 'geography' THEN NULL "
14561464
"ELSE 0 END AS DECIMAL_DIGITS, 10 AS NUM_PREC_RADIX, "
14571465
"NULL AS NULLABLE, NULL AS REMARKS, NULL AS COLUMN_DEF, "
14581466
"CAST(CASE columntype_rep "
@@ -1494,6 +1502,7 @@ def __build_local_schema_columns_query(
14941502
"WHEN 'geometry' THEN -4 "
14951503
"WHEN 'super' THEN -16 "
14961504
"WHEN 'varbyte' THEN -4 "
1505+
"WHEN 'geography' THEN -4 "
14971506
"ELSE 1111 END AS SMALLINT) AS SQL_DATA_TYPE, "
14981507
"CAST(NULL AS SMALLINT) AS SQL_DATETIME_SUB, CASE "
14991508
"WHEN LEFT (columntype,7) = 'varchar' THEN regexp_substr (columntype,'[0-9]+',7)::INTEGER "
@@ -1591,6 +1600,7 @@ def __build_universal_schema_columns_query(
15911600
" WHEN 'geometry' THEN -4 "
15921601
" WHEN 'super' THEN -16 "
15931602
" WHEN 'varbyte' THEN -4 "
1603+
" WHEN 'geography' THEN -4 "
15941604
" ELSE 1111 END AS SMALLINT) AS DATA_TYPE,"
15951605
" COALESCE("
15961606
" domain_name,"
@@ -1645,6 +1655,7 @@ def __build_universal_schema_columns_query(
16451655
" WHEN 'geometry' THEN NULL"
16461656
" WHEN 'super' THEN NULL"
16471657
" WHEN 'varbyte' THEN NULL"
1658+
" WHEN 'geography' THEN NULL "
16481659
" ELSE {unknown_column_size}"
16491660
" END AS COLUMN_SIZE,"
16501661
" NULL AS BUFFER_LENGTH,"
@@ -1659,6 +1670,7 @@ def __build_universal_schema_columns_query(
16591670
" WHEN 'geometry' THEN NULL"
16601671
" WHEN 'super' THEN NULL"
16611672
" WHEN 'varbyte' THEN NULL"
1673+
" WHEN 'geography' THEN NULL "
16621674
" ELSE 0"
16631675
" END AS DECIMAL_DIGITS,"
16641676
" 10 AS NUM_PREC_RADIX,"
@@ -1706,6 +1718,7 @@ def __build_universal_schema_columns_query(
17061718
" WHEN 'geometry' THEN -4"
17071719
" WHEN 'super' THEN -16"
17081720
" WHEN 'varbyte' THEN -4"
1721+
" WHEN 'geography' THEN -4 "
17091722
" ELSE 1111 END AS SMALLINT) AS SQL_DATA_TYPE,"
17101723
" CAST(NULL AS SMALLINT) AS SQL_DATETIME_SUB,"
17111724
" CASE data_type"
@@ -1746,6 +1759,7 @@ def __build_universal_schema_columns_query(
17461759
" WHEN 'geometry' THEN NULL"
17471760
" WHEN 'super' THEN NULL"
17481761
" WHEN 'varbyte' THEN NULL"
1762+
" WHEN 'geography' THEN NULL "
17491763
" ELSE {unknown_column_size}"
17501764
" END AS CHAR_OCTET_LENGTH,"
17511765
" ordinal_position AS ORDINAL_POSITION,"
@@ -1824,6 +1838,7 @@ def __build_universal_all_schema_columns_query(
18241838
" WHEN 'geometry' THEN -4 "
18251839
" WHEN 'super' THEN -16 "
18261840
" WHEN 'varbyte' THEN -4 "
1841+
" WHEN 'geography' THEN -4 "
18271842
" ELSE 1111 END AS SMALLINT) AS DATA_TYPE, "
18281843
" CASE data_type "
18291844
" WHEN 'boolean' THEN 'bool' "
@@ -1876,6 +1891,7 @@ def __build_universal_all_schema_columns_query(
18761891
" WHEN 'geometry' THEN NULL "
18771892
" WHEN 'super' THEN NULL "
18781893
" WHEN 'varbyte' THEN NULL "
1894+
" WHEN 'geography' THEN NULL "
18791895
" ELSE 2147483647 "
18801896
" END AS COLUMN_SIZE, "
18811897
" NULL AS BUFFER_LENGTH, "
@@ -1890,6 +1906,7 @@ def __build_universal_all_schema_columns_query(
18901906
" WHEN 'geometry' THEN NULL "
18911907
" WHEN 'super' THEN NULL "
18921908
" WHEN 'varbyte' THEN NULL "
1909+
" WHEN 'geography' THEN NULL "
18931910
" ELSE 0 "
18941911
" END AS DECIMAL_DIGITS, "
18951912
" 10 AS NUM_PREC_RADIX, "
@@ -1937,6 +1954,7 @@ def __build_universal_all_schema_columns_query(
19371954
" WHEN 'geometry' THEN -4 "
19381955
" WHEN 'super' THEN -16 "
19391956
" WHEN 'varbyte' THEN -4 "
1957+
" WHEN 'geography' THEN -4 "
19401958
" ELSE 1111 END AS SMALLINT) AS SQL_DATA_TYPE, "
19411959
" CAST(NULL AS SMALLINT) AS SQL_DATETIME_SUB, "
19421960
" CASE data_type "
@@ -1977,6 +1995,7 @@ def __build_universal_all_schema_columns_query(
19771995
" WHEN 'geometry' THEN NULL "
19781996
" WHEN 'super' THEN NULL "
19791997
" WHEN 'varbyte' THEN NULL "
1998+
" WHEN 'geography' THEN NULL "
19801999
" ELSE 2147483647 "
19812000
" END AS CHAR_OCTET_LENGTH, "
19822001
" ordinal_position AS ORDINAL_POSITION, "

‎redshift_connector/utils/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
date_in,
2222
date_recv_binary,
2323
float_array_recv,
24+
geographyhex_recv,
2425
h_pack,
2526
h_unpack,
2627
i_pack,

‎redshift_connector/utils/type_utils.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import enum
22
import typing
3+
from binascii import hexlify
34
from codecs import decode as codecs_decode
45
from collections import defaultdict
56
from datetime import date
@@ -49,6 +50,7 @@
4950
DATE_ARRAY = 1182
5051
FLOAT = 701
5152
FLOAT_ARRAY = 1022
53+
GEOGRAPHY = 3001
5254
GEOMETRY = 3000
5355
GEOMETRYHEX = 3999
5456
INET = 869
@@ -564,6 +566,10 @@ def hexencoding_lookup_no_case(input_value: int) -> int:
564566
return ascii_invalid_value
565567

566568

569+
def geographyhex_recv(data: bytes, idx: int, length: int) -> str:
570+
return hexlify(data[idx : idx + length]).decode(_client_encoding)
571+
572+
567573
def geometryhex_recv(data: bytes, idx: int, length: int) -> str:
568574
error_flag: bool = False
569575
pointer: int = idx
@@ -661,6 +667,7 @@ def varbytehex_recv(data: bytes, idx: int, length: int) -> str:
661667
NUMERIC: (FC_BINARY, numeric_in_binary), # NUMERIC
662668
# 2275: (FC_BINARY, text_recv), # cstring
663669
# 2950: (FC_BINARY, uuid_recv), # uuid
670+
GEOGRAPHY: (FC_BINARY, geographyhex_recv), # GEOGRAPHY
664671
GEOMETRY: (FC_TEXT, text_recv), # GEOMETRY
665672
GEOMETRYHEX: (FC_TEXT, geometryhex_recv), # GEOMETRYHEX
666673
# 3802: (FC_TEXT, json_in), # jsonb

‎test/integration/datatype/_generate_test_datatype_tables.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class RedshiftDatatypes(Enum):
2323
"""
2424

2525
super = auto()
26+
geography = auto()
2627
geometry = auto()
2728
varbyte = auto()
2829

@@ -154,6 +155,21 @@ def list(cls) -> typing.List["RedshiftDatatypes"]:
154155
"ffffffffffffff",
155156
)
156157
],
158+
# TODO: re-enable
159+
# RedshiftDatatypes.geography.name: (
160+
# (
161+
# "ST_GeogFromText('POINT(1 1)')",
162+
# '0101000020E6100000000000000000F03F000000000000F03F'
163+
# ),
164+
# (
165+
# "ST_GeogFromText('POINT(2 2)')",
166+
# '0101000020E610000000000000000000400000000000000040'
167+
# ),
168+
# (
169+
# "ST_GeogFromText('SRID=4267;POINT(-77.0092 38.889588)')",
170+
# '0101000020AB100000E3C798BB964053C000750305DE714340'
171+
# )
172+
# )
157173
}
158174

159175

‎test/integration/datatype/test_datatypes.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,6 @@ def test_datatype_recv_support(db_kwargs, datatype, client_protocol):
8282
redshift_datatype_testcases.append((datatype, test_case))
8383

8484

85-
@pytest.mark.skip(reason="wip")
8685
@pytest.mark.parametrize("client_protocol", ClientProtocolVersion.list())
8786
@pytest.mark.parametrize("_input", redshift_datatype_testcases)
8887
def test_redshift_specific_recv_support(db_kwargs, _input, client_protocol):
@@ -96,10 +95,12 @@ def test_redshift_specific_recv_support(db_kwargs, _input, client_protocol):
9695
results: typing.Tuple = cursor.fetchall()
9796
assert len(results) == 1
9897
assert len(results[0]) == 1
99-
assert results[0][0] == exp_val
98+
if datatype in ("varbyte", "geography"):
99+
assert results[0][0].lower() == exp_val.lower()
100+
else:
101+
assert results[0][0] == exp_val
100102

101103

102-
@pytest.mark.skip(reason="manual")
103104
@pytest.mark.parametrize("client_protocol", ClientProtocolVersion.list())
104105
@pytest.mark.parametrize("_input", redshift_test_data[RedshiftDatatypes.varbyte.name])
105106
def test_redshift_varbyte_insert(db_kwargs, _input, client_protocol):
@@ -117,7 +118,6 @@ def test_redshift_varbyte_insert(db_kwargs, _input, client_protocol):
117118
assert results[0][1] == bytes(data, encoding="utf-8").hex()
118119

119120

120-
@pytest.mark.skip(reason="manual")
121121
@pytest.mark.parametrize("client_protocol", ClientProtocolVersion.list())
122122
@pytest.mark.parametrize(
123123
"_input",

0 commit comments

Comments
 (0)