fix: remove SQLAlchemy LegacyAPIWarning from unit tests

azmeuk · azmeuk · commit eb4013471f9e · 2023-09-01T17:28:11.000+02:00
LegacyAPIWarning: The Query.get() method is considered legacy as of the 1.x series of SQLAlchemy and becomes a legacy construct in 2.0. The method is now available as Session.get() (deprecated since: 2.0) (Background on SQLAlchemy 2.0 at: https://sqlalche.me/e/b8d9) grant_user = User.query.get(int(user_id))
diff --git a/authlib/integrations/flask_oauth2/resource_protector.py b/authlib/integrations/flask_oauth2/resource_protector.py
@@ -38,7 +38,7 @@ def authenticate_token(self, token_string):
         @app.route('/user')
         @require_oauth(['profile'])
         def user_profile():
-            user = User.query.get(current_token.user_id)
+            user = User.get(current_token.user_id)
             return jsonify(user.to_dict())
 
     """
@@ -77,7 +77,7 @@ def acquire(self, scopes=None):
             @app.route('/api/user')
             def user_api():
                 with require_oauth.acquire('profile') as token:
-                    user = User.query.get(token.user_id)
+                    user = User.get(token.user_id)
                     return jsonify(user.to_dict())
         """
         try:
diff --git a/authlib/oauth2/rfc6749/grants/authorization_code.py b/authlib/oauth2/rfc6749/grants/authorization_code.py
@@ -339,7 +339,7 @@ def authenticate_user(self, authorization_code):
         MUST implement this method in subclass, e.g.::
 
             def authenticate_user(self, authorization_code):
-                return User.query.get(authorization_code.user_id)
+                return User.get(authorization_code.user_id)
 
         :param authorization_code: AuthorizationCode object
         :return: user
diff --git a/authlib/oauth2/rfc6749/grants/refresh_token.py b/authlib/oauth2/rfc6749/grants/refresh_token.py
@@ -158,7 +158,7 @@ def authenticate_user(self, refresh_token):
         implement this method in subclass::
 
             def authenticate_user(self, credential):
-                return User.query.get(credential.user_id)
+                return User.get(credential.user_id)
 
         :param refresh_token: Token object
         :return: user
diff --git a/authlib/oauth2/rfc8628/device_code.py b/authlib/oauth2/rfc8628/device_code.py
@@ -150,7 +150,7 @@ def query_device_credential(self, device_code):
         Developers MUST implement it in subclass::
 
             def query_device_credential(self, device_code):
-                return DeviceCredential.query.get(device_code)
+                return DeviceCredential.get(device_code)
 
         :param device_code: a string represent the code.
         :return: DeviceCredential instance
@@ -168,7 +168,7 @@ def query_user_grant(self, user_code):
                     return None
 
                 user_id, allowed = data.split()
-                user = User.query.get(user_id)
+                user = User.get(user_id)
                 return user, bool(allowed)
 
         Note, user grant information is saved by verification endpoint.
diff --git a/tests/flask/test_oauth1/oauth1_server.py b/tests/flask/test_oauth1/oauth1_server.py
@@ -215,7 +215,7 @@ def authorize():
                 return 'error'
         user_id = request.form.get('user_id')
         if user_id:
-            grant_user = User.query.get(int(user_id))
+            grant_user = db.session.get(User, int(user_id))
         else:
             grant_user = None
         try:
diff --git a/tests/flask/test_oauth2/models.py b/tests/flask/test_oauth2/models.py
@@ -38,7 +38,7 @@ class AuthorizationCode(db.Model, OAuth2AuthorizationCodeMixin):
 
     @property
     def user(self):
-        return User.query.get(self.user_id)
+        return db.session.get(User, self.user_id)
 
 
 class Token(db.Model, OAuth2TokenMixin):
@@ -64,7 +64,7 @@ def delete_authorization_code(self, authorization_code):
         db.session.commit()
 
     def authenticate_user(self, authorization_code):
-        return User.query.get(authorization_code.user_id)
+        return db.session.get(User, authorization_code.user_id)
 
 
 def save_authorization_code(code, request):
diff --git a/tests/flask/test_oauth2/oauth2_server.py b/tests/flask/test_oauth2/oauth2_server.py
@@ -36,7 +36,7 @@ def authorize():
         if request.method == 'GET':
             user_id = request.args.get('user_id')
             if user_id:
-                end_user = User.query.get(int(user_id))
+                end_user = db.session.get(User, int(user_id))
             else:
                 end_user = None
             try:
@@ -46,7 +46,7 @@ def authorize():
                 return url_encode(error.get_body())
         user_id = request.form.get('user_id')
         if user_id:
-            grant_user = User.query.get(int(user_id))
+            grant_user = db.session.get(User, int(user_id))
         else:
             grant_user = None
         return server.create_authorization_response(grant_user=grant_user)
diff --git a/tests/flask/test_oauth2/test_device_code_grant.py b/tests/flask/test_oauth2/test_device_code_grant.py
@@ -60,9 +60,9 @@ def query_device_credential(self, device_code):
 
     def query_user_grant(self, user_code):
         if user_code == 'code':
-            return User.query.get(1), True
+            return db.session.get(User, 1), True
         if user_code == 'denied':
-            return User.query.get(1), False
+            return db.session.get(User, 1), False
         return None
 
     def should_slow_down(self, credential):
diff --git a/tests/flask/test_oauth2/test_introspection_endpoint.py b/tests/flask/test_oauth2/test_introspection_endpoint.py
@@ -17,7 +17,7 @@ def query_token(self, token, token_type_hint):
         return query_token(token, token_type_hint)
 
     def introspect_token(self, token):
-        user = User.query.get(token.user_id)
+        user = db.session.get(User, token.user_id)
         return {
             "active": True,
             "client_id": token.client_id,
diff --git a/tests/flask/test_oauth2/test_refresh_token.py b/tests/flask/test_oauth2/test_refresh_token.py
@@ -15,7 +15,7 @@ def authenticate_refresh_token(self, refresh_token):
             return item
 
     def authenticate_user(self, credential):
-        return User.query.get(credential.user_id)
+        return db.session.get(User, credential.user_id)
 
     def revoke_old_credential(self, credential):
         now = int(time.time())
