Repository: kayak/pypika
Pull request #459 (2020-06-29 13:40:58)

Fix AtTimezone to call field get sql


Patch:
diff --git a/pypika/terms.py b/pypika/terms.py
index 02db605..b1ec0b3 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -467,7 +467,6 @@ class EmptyCriterion:
         return other
 
 
-
 class Field(Criterion, JSON):
     def __init__(self, name: str, alias: Optional[str] = None, table: Optional[Union[str, "Selectable"]] = None) -> None:
         super().__init__(alias)
@@ -1483,15 +1482,15 @@ class AtTimezone(Term):
     """
     is_aggregate = None
 
-    def __init__(self, name, zone, interval=False, alias=None):
+    def __init__(self, field, zone, interval=False, alias=None):
         super().__init__(alias)
-        self.name = name
+        self.field = Field(field) if not isinstance(field, Field) else field
         self.zone = zone
         self.interval = interval
 
     def get_sql(self, **kwargs):
         sql = '{name} AT TIME ZONE {interval}\'{zone}\''.format(
-            name=self.name,
+            name=self.field.get_sql(**kwargs),
             interval='INTERVAL ' if self.interval else '',
             zone=self.zone,
         )


Test patch:
diff --git a/pypika/tests/test_terms.py b/pypika/tests/test_terms.py
index 2831af3..9b83d8b 100644
--- a/pypika/tests/test_terms.py
+++ b/pypika/tests/test_terms.py
@@ -1,18 +1,46 @@
 from unittest import TestCase
 
-from pypika import Query
-from pypika.terms import AtTimezone
+from pypika import Query, Table
+from pypika.terms import AtTimezone, Field
 
 
 class AtTimezoneTests(TestCase):
     def test_when_interval_not_specified(self):
-        query = Query.from_('customers').select(AtTimezone('date', 'US/Eastern'))
-        self.assertEqual('SELECT date AT TIME ZONE \'US/Eastern\' FROM "customers"', str(query))
+        query = Query.from_("customers").select(AtTimezone("date", "US/Eastern"))
+        self.assertEqual(
+            'SELECT "date" AT TIME ZONE \'US/Eastern\' FROM "customers"', str(query)
+        )
 
     def test_when_interval_specified(self):
-        query = Query.from_('customers').select(AtTimezone('date', '-06:00', interval=True))
-        self.assertEqual('SELECT date AT TIME ZONE INTERVAL \'-06:00\' FROM "customers"', str(query))
+        query = Query.from_("customers").select(
+            AtTimezone("date", "-06:00", interval=True)
+        )
+        self.assertEqual(
+            'SELECT "date" AT TIME ZONE INTERVAL \'-06:00\' FROM "customers"',
+            str(query),
+        )
 
     def test_when_alias_specified(self):
-        query = Query.from_('customers').select(AtTimezone('date', 'US/Eastern', alias='alias1'))
-        self.assertEqual('SELECT date AT TIME ZONE \'US/Eastern\' "alias1" FROM "customers"', str(query))
+        query = Query.from_("customers").select(
+            AtTimezone("date", "US/Eastern", alias="alias1")
+        )
+        self.assertEqual(
+            'SELECT "date" AT TIME ZONE \'US/Eastern\' "alias1" FROM "customers"',
+            str(query),
+        )
+
+    def test_passes_kwargs_to_field_get_sql(self):
+        customers = Table("customers")
+        accounts = Table("accounts")
+        query = (
+            Query.from_(customers)
+            .join(accounts)
+            .on(customers.account_id == accounts.account_id)
+            .select(AtTimezone(customers.date, "US/Eastern", alias="alias1"))
+        )
+
+        self.assertEqual(
+            'SELECT "customers"."date" AT TIME ZONE \'US/Eastern\' "alias1" '
+            'FROM "customers" JOIN "accounts" ON "customers"."account_id"="accounts"."account_id"',
+            query.get_sql(with_namespace=True),
+        )