nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: kayak/pypika
Pull request #458 (2020-06-29 11:39:04)

Add AtTimezone support


Patch:
diff --git a/pypika/terms.py b/pypika/terms.py
index a6bf54b..02db605 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -1472,3 +1472,27 @@ class PseudoColumn(Term):
 
     def get_sql(self, **kwargs: Any) -> str:
         return self.name
+
+
+class AtTimezone(Term):
+    """
+    Generates AT TIME ZONE SQL.
+    Examples:
+        AT TIME ZONE 'US/Eastern'
+        AT TIME ZONE INTERVAL '-06:00'
+    """
+    is_aggregate = None
+
+    def __init__(self, name, zone, interval=False, alias=None):
+        super().__init__(alias)
+        self.name = name
+        self.zone = zone
+        self.interval = interval
+
+    def get_sql(self, **kwargs):
+        sql = '{name} AT TIME ZONE {interval}\'{zone}\''.format(
+            name=self.name,
+            interval='INTERVAL ' if self.interval else '',
+            zone=self.zone,
+        )
+        return format_alias_sql(sql, self.alias, **kwargs)


Test patch:
diff --git a/pypika/tests/test_terms.py b/pypika/tests/test_terms.py
new file mode 100644
index 0000000..2831af3
--- /dev/null
+++ b/pypika/tests/test_terms.py
@@ -0,0 +1,18 @@
+from unittest import TestCase
+
+from pypika import Query
+from pypika.terms import AtTimezone
+
+
+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))
+
+    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))
+
+    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))
Repository: kayak/pypika
Pull request #477 (2020-08-18 13:44:43)

add for_update statement
support sql like `select * from t where id=1 for update`


Patch:
diff --git a/pypika/queries.py b/pypika/queries.py
index 8406967..81ca923 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -598,6 +598,7 @@ class QueryBuilder(Selectable, Term):
         self._values = []
         self._distinct = False
         self._ignore = False
+        self._for_update = False
 
         self._wheres = None
         self._prewheres = None
@@ -832,6 +833,10 @@ class QueryBuilder(Selectable, Term):
     def distinct(self) -> "QueryBuilder":
         self._distinct = True
 
+    @builder
+    def for_update(self) -> "QueryBuilder":
+        self._for_update = True
+
     @builder
     def ignore(self) -> "QueryBuilder":
         self._ignore = True
@@ -1265,6 +1270,9 @@ class QueryBuilder(Selectable, Term):
         if self._offset:
             querystring += self._offset_sql()
 
+        if self._for_update:
+            querystring += self._for_update_sql()
+
         if subquery:
             querystring = "({query})".format(query=querystring)
 
@@ -1293,6 +1301,14 @@ class QueryBuilder(Selectable, Term):
 
         return distinct
 
+    def _for_update_sql(self) -> str:
+        if self._for_update:
+            for_update = ' FOR UPDATE'
+        else:
+            for_update = ''
+
+        return for_update
+
     def _select_sql(self, **kwargs: Any) -> str:
         return "SELECT {distinct}{select}".format(
               distinct=self._distinct_sql(**kwargs),


Test patch:
diff --git a/pypika/tests/test_selects.py b/pypika/tests/test_selects.py
index c629585..3a589a7 100644
--- a/pypika/tests/test_selects.py
+++ b/pypika/tests/test_selects.py
@@ -294,6 +294,10 @@ class WhereTests(unittest.TestCase):
         q = self.t.select("*").where(self.t.foo == self.t.bar)
         self.assertEqual(q, q1)
 
+    def test_where_field_equals_for_update(self):
+        q = Query.from_(self.t).select("*").where(self.t.foo == self.t.bar).for_update()
+        self.assertEqual('SELECT * FROM "abc" WHERE "foo"="bar" FOR UPDATE', str(q))
+
     def test_where_field_equals_where(self):
         q = (
             Query.from_(self.t)